File size: 835 Bytes
69591a9 |
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 |
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
|