| import os |
| import numpy as np |
| from tqdm import tqdm |
| from torch.utils.data import DataLoader |
| from torch.utils.data.dataset import Dataset |
| from PIL import Image |
| import torchvision.transforms.functional as TF |
| import torchvision.transforms as tf |
| from PIL import Image, ImageFile |
| import random |
| import math |
| from model import * |
| import torch |
| |
| |
|
|
| ImageFile.LOAD_TRUNCATED_IMAGES = True |
|
|
|
|
| class base_dataset(Dataset): |
| def __init__(self, data_dir, img_size, transforms=False, crop=False): |
| imgs = sorted(os.listdir(data_dir + "/input")) |
| self.input_imgs = [os.path.join(data_dir + "/input", name) for name in imgs] |
| self.gt_imgs = [os.path.join(data_dir + "/gt", name) for name in imgs] |
| self.transforms = transforms |
| self.crop = crop |
| self.img_size = img_size |
|
|
| def __getitem__(self, index): |
| inp_img_path = self.input_imgs[index] |
| gt_img_path = self.gt_imgs[index] |
| inp_img = Image.open(inp_img_path).convert("RGB") |
| gt_img = Image.open(gt_img_path).convert("RGB") |
| if self.transforms: |
| inp_img = self.transforms(inp_img) |
| gt_img = self.transforms(gt_img) |
|
|
| if self.crop: |
| inp_img, gt_img = self.crop_image(inp_img, gt_img) |
|
|
| return inp_img, gt_img, inp_img_path |
|
|
| def __len__(self): |
| return len(self.gt_imgs) |
|
|
| def crop_image(self, inp_img, gt_img): |
| crop_h, crop_w = self.img_size |
| i, j, h, w = tf.RandomCrop.get_params( |
| inp_img, output_size=((crop_h, crop_w))) |
| inp_img = TF.crop(inp_img, i, j, h, w) |
| gt_img = TF.crop(gt_img, i, j, h, w) |
| inp_img = TF.to_tensor(inp_img) |
| gt_img = TF.to_tensor(gt_img) |
|
|
| return inp_img, gt_img |
|
|
|
|
| class random_scale_dataset(Dataset): |
| def __init__(self, data_dir, img_size, transforms=False, crop=False): |
| imgs = sorted(os.listdir(data_dir + "/input")) |
| self.input_imgs = [os.path.join(data_dir + "/input", name) for name in imgs] |
| self.gt_imgs = [os.path.join(data_dir + "/gt", name) for name in imgs] |
| self.transforms = transforms |
| self.crop = crop |
| self.img_size = img_size |
|
|
| def __getitem__(self, index): |
| inp_img_path = self.input_imgs[index] |
| gt_img_path = self.gt_imgs[index] |
| inp_img = Image.open(inp_img_path).convert("RGB") |
| gt_img = Image.open(gt_img_path).convert("RGB") |
|
|
| random_scale_factor = random.randrange(self.img_size[0] * 0.25, self.img_size[0], 8) |
| down_h = down_w = random_scale_factor |
|
|
| if self.transforms: |
| inp_img = self.transforms(inp_img) |
| gt_img = self.transforms(gt_img) |
| return inp_img, gt_img, down_h, down_w, inp_img_path |
|
|
| if self.crop: |
| inp_img, gt_img = self.crop_image(inp_img, gt_img) |
| return inp_img, gt_img, down_h, down_w, inp_img_path |
|
|
| def __len__(self): |
| return len(self.gt_imgs) |
|
|
| def crop_image(self, inp_img, gt_img): |
| crop_h, crop_w = self.img_size |
| i, j, h, w = tf.RandomCrop.get_params( |
| inp_img, output_size=((crop_h, crop_w))) |
| inp_img = TF.crop(inp_img, i, j, h, w) |
| gt_img = TF.crop(gt_img, i, j, h, w) |
| inp_img = TF.to_tensor(inp_img) |
| gt_img = TF.to_tensor(gt_img) |
|
|
| return inp_img, gt_img |
|
|
|
|
| def get_loader(data_dir, img_size, transforms, crop_flag, batch_size, num_workers, shuffle, random_flag=False, inference_flag=False): |
| if random_flag: |
| dataset = random_scale_dataset(data_dir, img_size, transforms, crop_flag) |
| dataloader = DataLoader(dataset, batch_size=batch_size, |
| shuffle=shuffle, num_workers=num_workers, pin_memory=True) |
| else: |
| dataset = base_dataset(data_dir, img_size, transforms, crop_flag) |
| dataloader = DataLoader(dataset, batch_size=batch_size, |
| shuffle=shuffle, num_workers=num_workers, pin_memory=True) |
| return dataloader |
|
|
