siamese/coco_gen.py

75 lines
1.7 KiB
Python
Raw Normal View History

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("-----")