Spaces:
Sleeping
Sleeping
File size: 2,929 Bytes
1d64201 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import json
import numpy as np
from datetime import datetime
version_number = "0.0.1"
class AnnotationBuilder:
#creates the base structure of the coco format
def __init__(self):
self.licenses = []
self.categories = [{"id": 0,
"name": "Wall",
"supercategory": "none"},
{ "id": 1,
"name": "Door",
"supercategory": "none"},
{"id": 2,
"name": "Room",
"supercategory": "none"},
{"id": 3,
"name": "Window",
"supercategory": "none"}
]
self.images = []
self.annotations = []
def set_info(self, description, data_source_name, data_source_url, data_source_creation_date):
self.info = [{"year":2025,
"version":version_number,
"description":description,
"contributor":data_source_name,
"url":data_source_url,
"date_created":data_source_creation_date.strftime("%Y-%m-%dT%H:%M:%S")}]
def add_license(self, license_name, license_url):
self.licenses.append({"id":len(self.licenses),
"url":license_url,
"name":license_name})
def add_image(self, filename, width, height):
id = len(self.images)
self.images.append({"id":id,
"width":width,
"height":height,
"file_name":filename, #filename should be the image path relative to the cocofile's path
"license":0,
"date_captured":datetime.now().strftime("%Y-%m-%dT%H:%M:%S")})
return id
def add_annotation(self, image_id, category_id, poly):
id = len(self.annotations)
segmentation = np.array(poly.exterior.coords).astype(int).ravel().tolist()[:-2]
x,y,x2,y2 = tuple(map(int, poly.bounds))
self.annotations.append({"id":id,
"image_id":image_id,
"category_id":category_id,
"segmentation":[segmentation],
"area":poly.area,
"bbox":[x,y,x2-x,y2-y],
"iscrowd":0})
return id, poly
def final_output(self):
return {"info":self.info, "licenses":self.licenses, "categories":self.categories, "images":self.images, "annotations":self.annotations}
def save_file(self, filepath):
coco_file = open(filepath,'w')
json.dump(self.final_output(),coco_file,indent=4)
coco_file.close() |