| import os | |
| import os.path as osp | |
| import torch | |
| import torch.utils.data as data | |
| import data.util as util | |
| import random | |
| import numpy as np | |
| from PIL import Image | |
| class imageTestDataset(data.Dataset): | |
| def __init__(self, opt): | |
| super(imageTestDataset, self).__init__() | |
| self.opt = opt | |
| self.half_N_frames = opt['N_frames'] // 2 | |
| self.data_path = opt['data_path'] | |
| self.bit_path = opt['bit_path'] | |
| self.txt_path = self.opt['txt_path'] | |
| self.num_image = self.opt['num_image'] | |
| with open(self.txt_path) as f: | |
| self.list_image = f.readlines() | |
| self.list_image = [line.strip('\n') for line in self.list_image] | |
| self.list_image.sort() | |
| self.list_image = self.list_image | |
| l = len(self.list_image) // (self.num_image + 1) | |
| self.image_list_gt = self.list_image | |
| def __getitem__(self, index): | |
| path_GT = self.image_list_gt[index] | |
| img_GT = util.read_img(None, osp.join(self.data_path, path_GT)) | |
| img_GT = img_GT[:, :, [2, 1, 0]] | |
| img_GT = torch.from_numpy(np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float().unsqueeze(0) | |
| img_GT = torch.nn.functional.interpolate(img_GT, size=(512, 512), mode='nearest', align_corners=None) | |
| T, C, W, H = img_GT.shape | |
| list_h = [] | |
| R = 0 | |
| G = 0 | |
| B = 255 | |
| image = Image.new('RGB', (W, H), (R, G, B)) | |
| result = np.array(image) / 255. | |
| expanded_matrix = np.expand_dims(result, axis=0) | |
| expanded_matrix = np.repeat(expanded_matrix, T, axis=0) | |
| imgs_LQ = torch.from_numpy(np.ascontiguousarray(expanded_matrix)).float() | |
| imgs_LQ = imgs_LQ.permute(0, 3, 1, 2) | |
| imgs_LQ = torch.nn.functional.interpolate(imgs_LQ, size=(W, H), mode='nearest', align_corners=None) | |
| list_h.append(imgs_LQ) | |
| list_h = torch.stack(list_h, dim=0) | |
| return { | |
| 'LQ': list_h, | |
| 'GT': img_GT | |
| } | |
| def __len__(self): | |
| return len(self.image_list_gt) | |