import os
from scipy.io import loadmat
import numpy as np
from PIL import Image
from tqdm import tqdm

data_directory = "/vol/data/histo_datasets/CoNIC/"
images = np.load(data_directory+"images.npy")
masks = np.load(data_directory+"labels.npy")
for i in tqdm(range(len(images))):
    j = 0
    while np.max(masks[i,:,:,0]) > 255:
        mask = np.where(masks[i,:,:,0] < 256, masks[i,:,:,0], 0).astype(np.uint8)
        masks[i,:,:,0] = np.where(masks[i,:,:,0] > 255, masks[i,:,:,0] - 256, 0)
        Image.fromarray(mask).save(data_directory + "labels_png/" + str(i).zfill(4) + "_" + str(j) + ".png")
        j += 1
    Image.fromarray(images[i]).save(data_directory + "images_png/" + str(i).zfill(4) + ".png")
    Image.fromarray(masks[i,:,:,0]).save(data_directory + "labels_png/" + str(i).zfill(4) + ".png")