ai-training/annotation-queue/annotation_queue_dto.py

import json
from datetime import datetime, timedelta
from enum import Enum
import msgpack


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():
        with open('classes.json', '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


annotation_classes = AnnotationClass.read_json()


class AnnotationStatus(Enum):
    Created = 10
    Edited = 20
    Validated = 30
    Deleted = 40

    def __str__(self):
        return self.name


class SourceEnum(Enum):
    AI = 0
    Manual = 1


class RoleEnum(Enum):
    Operator = 10
    Validator = 20
    CompanionPC = 30
    Admin = 40
    ApiAdmin = 1000

    def __str__(self):
        return self.name

    def is_validator(self) -> bool:
        return self in {
            self.__class__.Validator,
            self.__class__.Admin,
            self.__class__.ApiAdmin
        }


class Detection:
    def __init__(self, annotation_name, cls, x, y, w, h, confidence=None):
        self.annotation_name = annotation_name
        self.cls = cls
        self.x = x
        self.y = y
        self.w = w
        self.h = h
        self.confidence = confidence

    def __str__(self):
        return f'{annotation_classes[self.cls].name}: {(self.confidence * 100):.1f}%'


class AnnotationCreatedMessageNarrow:
    def __init__(self, msgpack_bytes):
        unpacked_data = msgpack.unpackb(msgpack_bytes, strict_map_key=False)
        self.name = unpacked_data.get(1)
        self.createdEmail = unpacked_data.get(2)


class AnnotationMessage:
    def __init__(self, msgpack_bytes):
        unpacked_data = msgpack.unpackb(msgpack_bytes, strict_map_key=False)
        ts = unpacked_data[0]
        self.createdDate = datetime.utcfromtimestamp(ts.seconds) + timedelta(microseconds=ts.nanoseconds/1000)
        self.name = unpacked_data[1]
        self.originalMediaName = unpacked_data[2]
        self.time = timedelta(microseconds=unpacked_data[3]/10)
        self.imageExtension = unpacked_data[4]
        self.detections = self._parse_detections(unpacked_data[5])
        self.image = unpacked_data[6]

        self.createdRole = RoleEnum(unpacked_data[7])
        self.createdEmail = unpacked_data[8]
        self.source = SourceEnum(unpacked_data[9])
        self.status = AnnotationStatus(unpacked_data[10])

    def __str__(self):
        detections_str = ""
        if self.detections:
            detections_str_list = [str(detection) for detection in self.detections]
            detections_str = " [" + ", ".join(detections_str_list) + "]"
        createdBy = 'AI' if self.source == SourceEnum.AI else self.createdRole
        return f'{self.status} {self.name} by [{createdBy} {self.createdEmail}|{self.createdDate:%m-%d %H:%M}]{detections_str}'

    @staticmethod
    def _parse_detections(detections_json_str):
        if detections_json_str:
            detections_list = json.loads(detections_json_str)
            return [Detection(
                d.get('an'),
                d.get('cl'),
                d.get('x'),
                d.get('y'),
                d.get('w'),
                d.get('h'),
                d.get('p')
            ) for d in detections_list]
        return []


class AnnotationBulkMessage:
    def __init__(self, msgpack_bytes):
        unpacked_data = msgpack.unpackb(msgpack_bytes, strict_map_key=False)
        self.annotation_names = unpacked_data[0]
        self.annotation_status = AnnotationStatus(unpacked_data[1])
        self.createdEmail = unpacked_data[2]
        ts = unpacked_data[3]
        self.createdDate = datetime.utcfromtimestamp(ts.seconds) + timedelta(microseconds=ts.nanoseconds / 1000)

    def __str__(self):
        return f'{self.annotation_status}: [{self.annotation_names}] by [{self.createdEmail}|{self.createdDate:%m-%d %H:%M}]'