import json, os
import matplotlib.pyplot as plt
import numpy as np

from PIL import Image


no_label = 0
small = 0
passed = 0

count = 0


print("loading coco annotations...")
coco = json.load(open('./coco/annotations/instances_val2014.json'))
print("done")

def findAnnotationName(annotationId):
    for c in coco['categories']:
        if c['id'] == annotationId:
            return c['name']

def findAnnotationToId(ident):
    for annotation in coco['annotations']:
        img_an = annotation['image_id']
        if img_an == ident:
            return annotation 

def show(pil, pause=0.2):
    ImageNumpyFormat = np.asarray(pil)
    plt.imshow(ImageNumpyFormat)
    plt.draw()
    plt.pause(pause) # pause how many seconds
    plt.close()


def parseImage(coImg):
    # open image file

    path = "coco/val2014/" + coImg['file_name']
    img = Image.open(path)

    an = findAnnotationToId(coImg['id'])
    if an == None:
        no_label += 1
        return


    path = "coco/val2014/" + coImg['file_name']
    img = Image.open(path)
    c = an['bbox']
    crop = img.crop((c[0], c[1], c[0]+c[2], c[1]+c[3]))

    if crop.width < 64 or crop.height < 64:
        small += 1
        return

    imagePath = f"classified/{findAnnotationName(an['category_id'])}/{an['id']}.png"
    os.makedirs(os.path.dirname(imagePath), exist_ok=True)
    crop.save(imagePath)
    passed += 1

if __name__ == "__main__":
    for coImg in coco['images']:
        parseImage(coImg)
        count += 1
        if count % 100 == 0:
            print("status:")
            print(f"no labels: {no_label}")
            print(f"to small: {small")
            print(f"passed: {passed}")
            print("-----")