import os
from PIL import Image
import numpy as np
from tqdm import tqdm

data_source = "/vol/data/histo_datasets/PanNuke/"
folds = [("Fold 1/", "fold1/"), ("Fold 2/", "fold2/"), ("Fold 3/", "fold3/")]
image_counter = 0
mask_counter = 0
for f in folds:
    images = np.load(data_source + f[0] + "images/" + f[1] + "images.npy")
    for i in tqdm(images):
        Image.fromarray(i.astype(np.uint8)).save(data_source + "images_png/" + str(image_counter).zfill(4) + ".png")
        image_counter += 1
    masks = np.load(data_source + f[0] + "masks/" + f[1] + "masks.npy")
    for m in tqdm(masks):
        output = np.zeros((m.shape[0], m.shape[1]), dtype=np.int32)
        j = 0
        for i in range(m.shape[2]-1):
            values = np.unique(m[:, :, i])
            for v in values:
                if v != 0:
                    output[m[:, :, i] == v] = j
                    j += 1
        Image.fromarray(output).save(data_source + "masks_png/" + str(mask_counter).zfill(4) + ".png")
        mask_counter += 1