import cv2
import numpy as np
from skimage.segmentation import expand_labels

def read_to_mask(f):
    img = cv2.imread(str(f), cv2.IMREAD_UNCHANGED)[:,:,::-1]
    mask = np.zeros(img.shape[:2], dtype=np.uint8)
    mask[img[:, :, 0] > 200] = 1
    mask[img[:, :, 1] > 200] = 2

    return mask


def read_mask_from_path_gens(dict_gens, mapping=None):
    output = {k: dict() for k in dict_gens.keys()}
    for k, files in dict_gens.items():
        for file in files:
            name = file.parent.stem
            if mapping is not None:
                name = mapping.get(name, name)
            mask = read_to_mask(file)
            mask = expand_labels(mask, 1)
            if output[k].get(name) is None:
                output[k][name] = []
            output[k][name].append(mask)
    return output