import json from enum import Enum from os.path import join, dirname class Detection: def __init__(self, x, y, w, h, cls, confidence): self.x = x self.y = y self.w = w self.h = h self.cls = cls self.confidence = confidence def overlaps(self, det2, iou_threshold): overlap_x = 0.5 * (self.w + det2.w) - abs(self.x - det2.x) overlap_y = 0.5 * (self.h + det2.h) - abs(self.y - det2.y) intersection = max(0, overlap_x) * max(0, overlap_y) union = self.w * self.h + det2.w * det2.h - intersection return intersection / union > iou_threshold class Annotation: def __init__(self, frame, time, detections: list[Detection]): self.frame = frame self.time = time self.detections = detections if detections is not None else [] class WeatherMode(Enum): Norm = 0 Wint = 20 Night = 40 class AnnotationClass: def __init__(self, id, name, color): self.id = id self.name = name self.color = color color_str = color.lstrip('#') self.opencv_color = (int(color_str[4:6], 16), int(color_str[2:4], 16), int(color_str[0:2], 16)) @staticmethod def read_json(): classes_path = join(dirname(dirname(__file__)), 'classes.json') with open(classes_path, 'r', encoding='utf-8') as f: j = json.loads(f.read()) annotations_dict = {} for mode in WeatherMode: for cl in j: id = mode.value + cl['Id'] name = cl['Name'] if mode.value == 0 else f'{cl["Name"]}({mode.name})' annotations_dict[id] = AnnotationClass(id, name, cl['Color']) return annotations_dict @property def color_tuple(self): color = self.color[3:] lv = len(color) xx = range(0, lv, lv // 3) return tuple(int(color[i:i + lv // 3], 16) for i in xx)