ai-training/inference/dto.py

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)