import yaml
from PIL import Image
from data.base_dataset import BaseDataset

class ACDC(BaseDataset):
    def __init__(self, path_list, config_file = "./dataset/ACDC/acdc.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("/rgb_anon/", "/gt/")
            .replace("_rgb_", "_gt_")
            .replace("_anon.", "_labelIds.")
            for path in self.imgs
        ]