105 lines
4.2 KiB
Python
105 lines
4.2 KiB
Python
import os
|
|
import bisect
|
|
import threading
|
|
import torch
|
|
import numpy as np
|
|
import numpy.random as random
|
|
from PIL import Image
|
|
from torch.utils.data import Dataset
|
|
from folder_loader import FolderLoader
|
|
import torchvision.transforms as transforms
|
|
from utils import batch_edge_compute
|
|
|
|
def pil_loader(img_path):
|
|
return Image.open(img_path).convert("RGB")
|
|
|
|
|
|
class ImagePairPrefixFolder(Dataset):
|
|
def __init__(self, input_folder, gt_folder, max_img_size=0, size_unit=1, force_rgb=False):
|
|
super(ImagePairPrefixFolder, self).__init__()
|
|
|
|
self.gt_loader = FolderLoader(gt_folder)
|
|
# build the map from image name to index
|
|
self.gt_map = dict()
|
|
for idx, img_name in enumerate(self.gt_loader.img_names):
|
|
self.gt_map[os.path.splitext(img_name)[0].split('_')[0]] = idx
|
|
|
|
self.input_loader = FolderLoader(input_folder)
|
|
assert all([os.path.splitext(x)[0].split('_')[0] in self.gt_map for x in self.input_loader.img_names]), \
|
|
'cannot find corresponding gt names'
|
|
|
|
|
|
self.input_folder = input_folder
|
|
self.gt_folder = gt_folder
|
|
self.max_img_size = max_img_size
|
|
self.size_unit = size_unit
|
|
self.force_rgb = force_rgb
|
|
|
|
def __getitem__(self, index):
|
|
input_name, input_img = self.input_loader[index]
|
|
input_basename = os.path.splitext(input_name)[0].split('_')[0]
|
|
gt_idx = self.gt_map[input_basename]
|
|
|
|
gt_name, gt_img = self.gt_loader[gt_idx]
|
|
if self.force_rgb:
|
|
input_img = input_img.convert('RGB')
|
|
gt_img = gt_img.convert('RGB')
|
|
im_w, im_h = input_img.size
|
|
gt_w, gt_h = gt_img.size
|
|
assert im_w==gt_w and im_h==gt_h, 'input image and gt image size not match'
|
|
|
|
im_w, im_h = input_img.size
|
|
if 0 < self.max_img_size < max(im_w, im_h):
|
|
if im_w < im_h:
|
|
out_h = int(self.max_img_size) // self.size_unit * self.size_unit
|
|
out_w = int(im_w / im_h * out_h) // self.size_unit * self.size_unit
|
|
else:
|
|
out_w = int(self.max_img_size) // self.size_unit * self.size_unit
|
|
out_h = int(im_h / im_w * out_w) // self.size_unit * self.size_unit
|
|
else:
|
|
out_w = im_w // self.size_unit * self.size_unit
|
|
out_h = im_h // self.size_unit * self.size_unit
|
|
|
|
if im_w != out_w or im_h != out_h:
|
|
input_img = input_img.resize((out_w, out_h), Image.BILINEAR)
|
|
gt_img = gt_img.resize((out_w, out_h), Image.BILINEAR)
|
|
|
|
im_w, im_h = input_img.size
|
|
|
|
input_img = np.array(input_img).astype('float')
|
|
gt_img = np.array(gt_img).astype('float')
|
|
if len(input_img.shape) == 2:
|
|
input_img = input_img[:, :, np.newaxis]
|
|
if len(gt_img.shape) == 2:
|
|
gt_img = gt_img[:, :, np.newaxis]
|
|
return {'input_img': input_img, 'gt_img': gt_img, 'input_h': im_h, "input_w": im_w}
|
|
|
|
def get_input_info(self, index):
|
|
image_name = os.path.splitext(self.input_loader.img_names[index])[0]
|
|
return self.input_loader, image_name
|
|
|
|
def __len__(self):
|
|
return len(self.input_loader)
|
|
|
|
|
|
def var_custom_collate(batch):
|
|
min_h, min_w = 10000, 10000
|
|
for item in batch:
|
|
min_h = min(min_h, item['input_h'])
|
|
min_w = min(min_w, item['input_w'])
|
|
inc = 1 if len(batch[0]['input_img'].shape)==2 else batch[0]['input_img'].shape[2]
|
|
batch_input_images = torch.Tensor(len(batch), inc, min_h, min_w)
|
|
batch_gt_images = torch.Tensor(len(batch), inc, min_h, min_w)
|
|
|
|
for idx, item in enumerate(batch):
|
|
off_y = 0 if item['input_h']==min_h else random.randint(0, item['input_h'] - min_h)
|
|
off_x = 0 if item['input_w']==min_w else random.randint(0, item['input_w'] - min_w)
|
|
crop_input_img = item['input_img'][off_y:off_y + min_h, off_x:off_x + min_w, :]
|
|
crop_gt_img = item['gt_img'][off_y:off_y + min_h, off_x:off_x + min_w, :]
|
|
batch_input_images[idx] = torch.from_numpy(crop_input_img.transpose((2, 0, 1))) - 128
|
|
batch_gt_images[idx] = torch.from_numpy(crop_gt_img.transpose((2, 0, 1)))
|
|
|
|
|
|
batch_input_edges = batch_edge_compute(batch_input_images) - 128
|
|
return batch_input_images, batch_input_edges, batch_gt_images
|