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




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

def findAnnotationToId(ident, annotations):
    for annotation in annotations['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, annotations, subset):
    global no_label, small, passed
    # open image file

    path = "coco/"+subset+"/images/" + coImg['file_name'].split('/')[4]
    img = Image.open(path)

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

    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/classified_{subset}/{findAnnotationName(an['category_id'], annotations)}/{an['id']}.png"
    os.makedirs(os.path.dirname(imagePath), exist_ok=True)
    crop.save(imagePath)
    passed += 1

def parseSubset(subset):
    global count

    print("loading" + subset + "annotations...")
    annotations = json.load(open('./coco/'+subset+'/annotations/instances_default.json'))
    print("done")

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


if __name__ == "__main__":
    parseSubset('donaulager')
    parseSubset('instantina')
    parseSubset('lkw_walter_2019')
    parseSubset('lkw_walter_2020')
    parseSubset('tkl')
    parseSubset('vw_bratislava')
    parseSubset('vw_portugal')
    parseSubset('watt')
    parseSubset('wls')