CellPilot / preprocessing_scripts /preprocess_segpc.py
philippendres's picture
Upload folder using huggingface_hub
907462b verified
Raw
History Blame Contribute Delete
1.18 kB
import numpy as np
import os
from PIL import Image
data_source = "/vol/data/histo_datasets/SegPC/TCIA_SegPC_dataset/"
splits = ["train", "validation"]
for s in splits:
images_list = os.listdir(data_source + s + "/x/")
masks_list = os.listdir(data_source + s + "/y/")
for i in images_list:
short_masks_list = [m for m in masks_list if i[:-4] == m[:-6]]
j = 0
for m in short_masks_list:
if j == 0:
mask = np.array(Image.open(data_source + s + "/y/" + m))
if len(mask.shape) == 3:
mask = mask[:, :, 0]
mask = np.where(mask == 20, 1, mask)
mask = np.where(mask == 40, 2, mask)
else:
additional_mask = np.array(Image.open(data_source + s + "/y/" + m))
if len(additional_mask.shape) == 3:
additional_mask = additional_mask[:, :, 0]
mask = np.where(additional_mask == 20, mask + 2*j -1, mask)
mask = np.where(additional_mask == 40, mask + 2*j, mask)
j += 1
Image.fromarray(mask).save(data_source + s + "/masks_png/" + m[:-6] + ".png")