Spaces:
Runtime error
Runtime error
File size: 1,803 Bytes
24be7a2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import os
import os.path
import random
import numpy as np
import torch
import torch.utils.data as data
from PIL import Image
class MaskDataset(data.Dataset):
def __init__(self, segm_dir, ann_dir, downsample_factor=2, xflip=False):
self._segm_path = segm_dir
self._image_fnames = []
self.downsample_factor = downsample_factor
self.xflip = xflip
# load attributes
assert os.path.exists(f'{ann_dir}/upper_fused.txt')
for idx, row in enumerate(
open(os.path.join(f'{ann_dir}/upper_fused.txt'), 'r')):
annotations = row.split()
self._image_fnames.append(annotations[0])
def _open_file(self, path_prefix, fname):
return open(os.path.join(path_prefix, fname), 'rb')
def _load_segm(self, raw_idx):
fname = self._image_fnames[raw_idx]
fname = f'{fname[:-4]}_segm.png'
with self._open_file(self._segm_path, fname) as f:
segm = Image.open(f)
if self.downsample_factor != 1:
width, height = segm.size
width = width // self.downsample_factor
height = height // self.downsample_factor
segm = segm.resize(
size=(width, height), resample=Image.NEAREST)
segm = np.array(segm)
# segm = segm[:, :, np.newaxis].transpose(2, 0, 1)
return segm.astype(np.float32)
def __getitem__(self, index):
segm = self._load_segm(index)
if self.xflip and random.random() > 0.5:
segm = segm[:, ::-1].copy()
segm = torch.from_numpy(segm).long()
return_dict = {'segm': segm, 'img_name': self._image_fnames[index]}
return return_dict
def __len__(self):
return len(self._image_fnames)
|