| from torch.utils import data |
| import os |
| from PIL import Image, ImageFile |
|
|
|
|
| ImageFile.LOAD_TRUNCATED_IMAGES = True |
|
|
|
|
| class EvalDataset(data.Dataset): |
| def __init__(self, pred_root, label_root, return_predpath=False, return_gtpath=False): |
| self.return_predpath = return_predpath |
| self.return_gtpath = return_gtpath |
| pred_dirs = os.listdir(pred_root) |
| label_dirs = os.listdir(label_root) |
|
|
| dir_name_list = [] |
| for idir in pred_dirs: |
| if idir in label_dirs: |
| pred_names = os.listdir(os.path.join(pred_root, idir)) |
| label_names = os.listdir(os.path.join(label_root, idir)) |
| for iname in pred_names: |
| if iname in label_names: |
| dir_name_list.append(os.path.join(idir, iname)) |
|
|
| self.image_path = list( |
| map(lambda x: os.path.join(pred_root, x), dir_name_list)) |
| self.label_path = list( |
| map(lambda x: os.path.join(label_root, x), dir_name_list)) |
|
|
| self.labels = [] |
| for p in self.label_path: |
| self.labels.append(Image.open(p).convert('L')) |
|
|
|
|
| def __getitem__(self, item): |
| predpath = self.image_path[item] |
| gtpath = self.label_path[item] |
| pred = Image.open(predpath).convert('L') |
| gt = self.labels[item] |
| if pred.size != gt.size: |
| pred = pred.resize(gt.size, Image.BILINEAR) |
| returns = [pred, gt] |
| if self.return_predpath: |
| returns.append(predpath) |
| if self.return_gtpath: |
| returns.append(gtpath) |
| return returns |
|
|
| def __len__(self): |
| return len(self.image_path) |
|
|
