|
import os
|
|
import cv2
|
|
import math
|
|
import numpy as np
|
|
from torch.utils.data import Dataset
|
|
import os.path
|
|
import random
|
|
import torchvision.transforms as transforms
|
|
import torch
|
|
from PIL import Image, ImageDraw
|
|
|
|
|
|
class Image_Editing_Dataset(Dataset):
|
|
def __init__(self, cfg, dataset_root, split='test', dataset_name=''):
|
|
self.split = split
|
|
self.cfg = cfg
|
|
self.dataset_name = dataset_name
|
|
self.img_format = '.png'
|
|
|
|
self.dir_img = os.path.join(dataset_root, 'test_processed', 'images')
|
|
self.dir_lab = os.path.join(dataset_root, 'test_processed', 'labels')
|
|
self.dir_ins = os.path.join(dataset_root, 'test_processed', 'inst_map')
|
|
name_list = os.listdir(self.dir_img)
|
|
self.name_list = [n[:-4] for n in name_list if n.endswith(self.img_format)]
|
|
self.name_list.sort()
|
|
self.predefined_mask_path = os.path.join(dataset_root, f'test_processed', 'predefined_masks')
|
|
|
|
def __getitem__(self, index):
|
|
name = self.name_list[index]
|
|
|
|
img = cv2.imread(os.path.join(self.dir_img, name + '.png'))
|
|
lab = cv2.imread(os.path.join(self.dir_lab, name + '.png'), 0)
|
|
inst_map = Image.open(os.path.join(self.dir_ins, name + '.png'))
|
|
inst_map = np.array(inst_map, dtype=np.int32)
|
|
|
|
assert len(inst_map.shape) == 2
|
|
|
|
img = get_transform(img)
|
|
lab = get_transform(lab, normalize=False)
|
|
lab = lab * 255.0
|
|
|
|
mask = cv2.imread(os.path.join(self.predefined_mask_path, 'type_0', name + '.png'), 0) / 255
|
|
mask = mask.reshape((1,) + mask.shape).astype(np.float32)
|
|
|
|
mask = torch.from_numpy(mask)
|
|
masked_img = img * (1. - mask)
|
|
|
|
inst_map = inst_map.reshape((1,) + inst_map.shape).astype(np.float32)
|
|
inst_map = torch.from_numpy(inst_map)
|
|
|
|
return {'img': img, 'masked_img': masked_img, 'lab': lab, 'mask': mask, 'inst_map': inst_map, 'name': name}
|
|
|
|
|
|
def __len__(self):
|
|
"""Return the total number of images in the dataset."""
|
|
return len(self.name_list)
|
|
|
|
|
|
def get_transform(img, normalize=True):
|
|
transform_list = []
|
|
|
|
transform_list += [transforms.ToTensor()]
|
|
if normalize:
|
|
transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
|
|
return transforms.Compose(transform_list)(img)
|
|
|
