73 lines
1.7 KiB
Python
73 lines
1.7 KiB
Python
|
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_default.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):
|
||
|
|
global no_label, small, passed, coco
|
||
|
|
# open image file
|
||
|
|
|
||
|
|
path = "coco/roto_frank/images/" + coImg['file_name'].split('/')[4]
|
||
|
|
img = Image.open(path)
|
||
|
|
|
||
|
|
an = findAnnotationToId(coImg['id'])
|
||
|
|
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_roto/{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("-----")
|
||
|
|
|