from pycocotools.coco import COCO
import os
from PIL import Image
import numpy as np
from matplotlib import pyplot as plt

data_directory = "/vol/data/histo_datasets/CRAG/"
datasets = ["train2017", "val2017"]
for d in datasets:
    coco = COCO(data_directory+'cell_CRAG/annotations/instances_' + d + '.json')
    cat_ids = coco.getCatIds()
    for image_id in range(len(coco.imgs)):
        img = coco.imgs[image_id+1]
        if "aug" in img['file_name']:
            continue  
        anns_ids = coco.getAnnIds(imgIds=image_id+1, catIds=cat_ids, iscrowd=None)
        anns = coco.loadAnns(anns_ids)
        mask = coco.annToMask(anns[0])
        for i in range(len(anns)):
            mask += (i+2)*coco.annToMask(anns[i])
        mask = mask.astype(np.uint8)
        Image.fromarray(mask).save(data_directory + "cell_CRAG/labels/" + img['file_name'])