Convert bounding boxes from YOLO format into JSON format expected by edge-impulse-uploader

convert.py

import os, glob, json
import cv2

class BoundingBoxConvertor:
	def __init__(self):
		self.label_map = {
			"0":"class1",
			"1":"class2",
			"2":"class3"
		}

	def _image_annotation_object(self, annotation_file, image_file):
		print(image_file)

		with open(annotation_file) as infile:
			txt_annotations = infile.read().splitlines()

		height, width, _ = cv2.imread(image_file).shape

		annotations = []
		for each_annotation in txt_annotations:
			tokens = each_annotation.split(' ')

			annotations.append(
				{
					"label": self.label_map[tokens[0]],
					"x": float(tokens[1]) * width - (float(tokens[3]) * width)/2,
					"y": float(tokens[2]) * height - (float(tokens[4]) * height)/2,
					"width": float(tokens[3]) * width,
					"height": float(tokens[4]) * height
				}
			)
		return annotations


	def convert(self, directory, image_file_extension = ".png"):
		if not os.path.isdir(directory):
			raise Exception(f"Directory {directory} does not exist")

		annotation_file_paths = glob.glob(f"{directory}/*.txt")
		image_annotation_objects = {}

		for annotation_file_path in annotation_file_paths:
			image_file_path = os.path.splitext(annotation_file_path)[0] + image_file_extension
			image_file_name = os.path.basename(image_file_path)
			image_annotation_objects[image_file_name] = self._image_annotation_object(annotation_file_path, image_file_path)

		with open(f"{directory}/bounding_boxes.labels", 'w') as outfile:
			json.dump(
				{
					"version": 1,
					"type": "bounding-box-labels",
					"boundingBoxes": image_annotation_objects
				},
				outfile
			)


	
if __name__ == "__main__": 
	converter = BoundingBoxConvertor()
	converter.convert(
		# Directory containing images and text annotation files
		directory="dataset",
		# File extension of images in the directory
		image_file_extension = ".png"
	)