def aug(self, img, label):
# img, label = colorjitter(img, label)
# img, label = random_resize(img, label)
# img, label = random_rotate(img, label)s
img, label = random_crop(img, label, self.crop_size)
img, label = random_flip(img, label)
return img, label
def __getitem__(self, index):
index = index % self.length
lr_img = []
hr_img = []
for i in range(self.nframes):
lr_img.append(cv2.imread(self.imglist[index][0][i]))
lr_img = np.stack(lr_img, -1)
hr_img = cv2.imread(self.imglist[index][1][self.nframes - 1])
lr_img, hr_img = random_crop_new(lr_img, hr_img, self.crop_size)
lr_img, hr_img = random_flip(lr_img, hr_img)
lr_target = lr_img[:, :, :, -1]
bicubic = cv2.resize(lr_target,
dsize=None,
fx=4,
fy=4,
interpolation=cv2.INTER_CUBIC)
flow = [
get_flow(lr_target, lr_img[:, :, :, j])
for j in range(self.nframes - 1)
]
lr_target = lr_target.transpose(2, 0, 1) / 255
hr_img = hr_img.transpose(2, 0, 1) / 255
neighbor = [
lr_img[:, :, :, j].transpose(2, 0, 1) / 255
for j in range(self.nframes - 1)
]
flow = [j.transpose(2, 0, 1) for j in flow]
bicubic = bicubic.transpose(2, 0, 1) / 255
return lr_target, hr_img, neighbor, flow, bicubic
def __getitem__(self,index):
index = index % self.length
lr_img = []
hr_img = []
for i in range(self.nframes):
lr_img.append(cv2.imread(self.imglist[index][0][i]))
hr_img.append(cv2.imread(self.imglist[index][1][i]))
lr_img = np.stack(lr_img, -1)
hr_img = np.stack(hr_img, -1)
lr_img, hr_img = random_crop_new(lr_img, hr_img, self.crop_size)
lr_img, hr_img = random_flip(lr_img, hr_img)
lr_img = lr_img.transpose(3, 2, 0, 1) / 255
hr_img = hr_img.transpose(3, 2, 0, 1) / 255
hr_img = hr_img[self.nframes//2]
return lr_img, hr_img
def aug(self, img, label):
img, label = random_crop_new(img, label, self.crop_size, 0, 0)
img, label = random_flip(img, label)
return img, label