| 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 | |