| import os | |
| import torch.utils.data as data | |
| import torchvision.transforms as T | |
| from PIL import Image | |
| class LowLightDatasetTest(data.Dataset): | |
| def __init__(self, root, reside=False): | |
| self.root = root | |
| self.items = [] | |
| subsets = os.listdir(root) | |
| for subset in subsets: | |
| img_root = os.path.join(root, subset) | |
| img_names = list(sorted(os.listdir(img_root))) | |
| for img_name in img_names: | |
| self.items.append(( | |
| os.path.join(img_root, img_name), | |
| subset, | |
| img_name | |
| )) | |
| self.preproc = T.Compose( | |
| [T.ToTensor()] | |
| ) | |
| self.preproc_raw = T.Compose( | |
| [T.ToTensor()] | |
| ) | |
| def __getitem__(self, idx): | |
| img_path, subset, img_name = self.items[idx] | |
| img = Image.open(img_path).convert("RGB") | |
| img = img.resize((img.width // 8 * 8, img.height // 8 * 8), Image.ANTIALIAS) | |
| img_raw = self.preproc_raw(img) | |
| return img_raw, subset, img_name | |
| def __len__(self): | |
| return len(self.items) | |