Fix bug in 'corrupt_before_downsize=true'

This commit is contained in:
James Betker 2020-12-16 09:41:59 -07:00
parent 8661207d57
commit 4310e66848

View File

@ -4,6 +4,7 @@ from data.image_corruptor import ImageCorruptor
from data.chunk_with_reference import ChunkWithReference from data.chunk_with_reference import ChunkWithReference
import os import os
import cv2 import cv2
import numpy as np
# Class whose purpose is to hold as much logic as can possibly be shared between datasets that operate on raw image # Class whose purpose is to hold as much logic as can possibly be shared between datasets that operate on raw image
# data and nothing else (which also have a very specific directory structure being used, as dictated by # data and nothing else (which also have a very specific directory structure being used, as dictated by
@ -101,7 +102,7 @@ class BaseUnsupervisedImageDataset(data.Dataset):
h, w, _ = hs[0].shape h, w, _ = hs[0].shape
ls, lrs, lms, lcs = [], [], [], [] ls, lrs, lms, lcs = [], [], [], []
if self.corrupt_before_downsize and not self.for_eval: if self.corrupt_before_downsize and not self.for_eval:
hs = self.corruptor.corrupt_images(hs) hs = self.corruptor.corrupt_images(np.copy(hs))
for hq, hq_ref, hq_mask, hq_center in zip(hs, hrefs, hmasks, hcenters): for hq, hq_ref, hq_mask, hq_center in zip(hs, hrefs, hmasks, hcenters):
if self.for_eval: if self.for_eval:
ls.append(hq) ls.append(hq)