add rabbit consumer

This commit is contained in:
Alex Bezdieniezhnykh
2025-03-05 20:58:33 +02:00
parent 2fa864018f
commit b4a4826b19
6 changed files with 152 additions and 3 deletions
+87
View File
@@ -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}")