File size: 749 Bytes
e8f2edc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
import yaml
from PIL import Image
from data.base_dataset import BaseDataset
class CADEdgeTune(BaseDataset):
def __init__(self, path_list, config_file = "./dataset/CADEdgeTune/cadedgetune.yaml", transform = None,
data_set = 'val', seed=None, img_size=768, interpolation=Image.BILINEAR, color_pallete = 'city'):
super().__init__(path_list, transform, data_set, seed, img_size, interpolation, color_pallete)
with open(config_file, 'r') as stream:
cityyaml = yaml.safe_load(stream)
self.learning_map = cityyaml['learning_map']
self.masks = [
path.replace("/Images/", "/LabelMasks/")
.replace(".jpg", "_label_mask.png")
for path in self.imgs
]
|