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