import msgpack from pathlib import Path cdef class Detection: def __init__(self, double x, double y, double w, double h, int cls, double confidence): self.annotation_name = None self.x = x self.y = y self.w = w self.h = h self.cls = cls self.confidence = confidence def __str__(self): return f'{self.cls}: {self.x:.2f} {self.y:.2f} {self.w:.2f} {self.h:.2f}, prob: {(self.confidence*100):.1f}%' cdef overlaps(self, Detection det2): cdef double overlap_x = 0.5 * (self.w + det2.w) - abs(self.x - det2.x) cdef double overlap_y = 0.5 * (self.h + det2.h) - abs(self.y - det2.y) cdef double overlap_area = max(0.0, overlap_x) * max(0.0, overlap_y) cdef double min_area = min(self.w * self.h, det2.w * det2.h) return overlap_area / min_area > 0.6 cdef class Annotation: def __init__(self, str name, long ms, list[Detection] detections): self.original_media_name = Path(name).stem.replace(" ", "") self.name = f'{self.original_media_name}_{self.format_time(ms)}' self.time = ms self.detections = detections if detections is not None else [] for d in self.detections: d.annotation_name = self.name self.image = b'' def __str__(self): if not self.detections: return f"{self.name}: No detections" detections_str = ", ".join( f"class: {d.cls} {d.confidence * 100:.1f}% ({d.x:.2f}, {d.y:.2f})" for d in self.detections ) return f"{self.name}: {detections_str}" cdef format_time(self, ms): # Calculate hours, minutes, seconds, and hundreds of milliseconds. h = ms // 3600000 # Total full hours. ms_remaining = ms % 3600000 m = ms_remaining // 60000 # Full minutes. ms_remaining %= 60000 s = ms_remaining // 1000 # Full seconds. f = (ms_remaining % 1000) // 100 # Hundreds of milliseconds. h = h % 10 return f"{h}{m:02}{s:02}{f}" cdef bytes serialize(self): return msgpack.packb({ "n": self.name, "mn": self.original_media_name, "i": self.image, # "i" = image "t": self.time, # "t" = time "d": [ # "d" = detections { "an": det.annotation_name, "x": det.x, "y": det.y, "w": det.w, "h": det.h, "c": det.cls, "p": det.confidence } for det in self.detections ] })