add rabbit consumer

2026-04-22 23:56:35 +00:00 · 2025-03-05 20:58:33 +02:00
parent 2fa864018f
commit b4a4826b19
6 changed files with 152 additions and 3 deletions
@@ -0,0 +1,87 @@
+import os
+import msgpack
+import json
+
+
+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'{self.cls}: {self.x:.2f} {self.y:.2f} {self.w:.2f} {self.h:.2f}, prob: {(self.confidence * 100):.1f}%'
+
+
+class AnnotationCreatedMessage:
+    def __init__(self, msgpack_bytes):
+        unpacked_data = msgpack.unpackb(msgpack_bytes, raw=False)
+        self.createdDate = unpacked_data.get(0)
+        self.name = unpacked_data.get(1)
+        self.originalMediaName = unpacked_data.get(2)
+        self.time = unpacked_data.get(3)
+        self.imageExtension = unpacked_data.get(4)
+        detections_json_str = unpacked_data.get(5)
+        self.detections = self._parse_detections(detections_json_str)
+        self.image = unpacked_data.get(6)
+        self.createdRole = unpacked_data.get(7)
+        self.createdEmail = unpacked_data.get(8)
+        self.source = unpacked_data.get(9)
+        self.status = unpacked_data.get(10)
+
+    def __str__(self):
+        if self.detections:
+            detections_str_list = [str(detection) for detection in self.detections]
+            detections_str = ", ".join(detections_str_list)
+            return f'{self.name}: [{detections_str}]'
+        else:
+            return f'{self.name}: [Empty]'
+
+    @staticmethod
+    def _parse_detections(detections_json_str):
+        if detections_json_str:
+            detections_list = json.loads(detections_json_str)
+            return [Detection(
+                d.get('AnnotationName'),
+                d.get('ClassNumber'),
+                d.get('CenterX'),
+                d.get('CenterY'),
+                d.get('Height'),
+                d.get('Width'),
+                d.get('Probability')
+            ) for d in detections_list]
+        return []
+
+    def save_annotation(self, save_folder):
+        image_folder = os.path.join(save_folder, 'images')
+        labels_folder = os.path.join(save_folder, 'labels')
+
+        os.makedirs(image_folder, exist_ok=True)
+        os.makedirs(labels_folder, exist_ok=True)
+
+        image_path = os.path.join(image_folder, f"{self.name}.{self.imageExtension}")
+        label_path = os.path.join(labels_folder, f"{self.name}.txt")
+
+        try:
+            with open(image_path, 'wb') as image_file:
+                image_file.write(self.image)
+            print(f"Image saved to: {image_path}")
+        except IOError as e:
+            print(f"Error saving image: {e}")
+
+        try:
+            with open(label_path, 'w') as label_file:
+                if self.detections:
+                    label_file.writelines([
+                        f'{detection.cls} {detection.x} {detection.y} {detection.w} {detection.h}'
+                        for detection in self.detections
+                    ])
+                else:
+                    label_file.write('')
+            print(f'Label saved to: {label_path}')
+        except IOError as e:
+            print(f"Error saving label: {e}")