|
|
|
|
|
import os |
|
from torch.utils.data import Dataset |
|
from PIL import Image |
|
import torch |
|
|
|
class GTResDataset(Dataset): |
|
|
|
def __init__(self, root_path, gt_dir=None, transform=None, transform_train=None): |
|
self.pairs = [] |
|
for f in os.listdir(root_path): |
|
image_path = os.path.join(root_path, f) |
|
gt_path = os.path.join(gt_dir, f) |
|
if f.endswith(".jpg") or f.endswith(".png"): |
|
self.pairs.append([image_path, gt_path.replace('.png', '.jpg'), None]) |
|
self.transform = transform |
|
self.transform_train = transform_train |
|
|
|
def __len__(self): |
|
return len(self.pairs) |
|
|
|
def __getitem__(self, index): |
|
from_path, to_path, _ = self.pairs[index] |
|
from_im = Image.open(from_path).convert('RGB') |
|
to_im = Image.open(to_path).convert('RGB') |
|
|
|
if self.transform: |
|
to_im = self.transform(to_im) |
|
from_im = self.transform(from_im) |
|
|
|
return from_im, to_im |
|
|
