azaion/ai-training
1
0
Fork 0
mirror of https://github.com/azaion/ai-training.git synced 2026-04-22 09:26:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix rabbit consumer

Browse Source
This commit is contained in:
Alex Bezdieniezhnykh
2025-03-07 14:33:02 +02:00
parent 4e8c67995a
commit edf685fabc
3 changed files with 101 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
annotation_queue_handler.py
+33 -22
View File
@@ -6,51 +6,62 @@ from dto.annotation_created_message import AnnotationCreatedMessage
from utils import Dotdict
async def main():
class AnnotationQueueHandler:
def __init__(self):
with open(constants.CONFIG_FILE, "r") as f:
config_dict = yaml.safe_load(f)
d_config = Dotdict(config_dict)
queue_config = Dotdict(d_config.queue)
try:
with open(constants.OFFSET_FILE, 'r') as f:
offset_dict = yaml.safe_load(f)
offset_queue = offset_dict['offset_queue']
offset_confirmed_queue = offset_dict['offset_confirmed_queue']
except (FileNotFoundError, ValueError):
with open(constants.OFFSET_FILE, 'w') as f:
f.writelines(['offset_queue: 0',
'offset_confirmed_queue: 0'])
offset_queue = 0
offset_confirmed_queue = 0
consumer = Consumer(
self.consumer = Consumer(
host=queue_config.host,
port=queue_config.port,
username=queue_config.consumer_user,
password=queue_config.consumer_pw
)
try:
with open(constants.OFFSET_FILE, 'r') as f:
offset_dict = yaml.safe_load(f)
self.offset_queue = offset_dict['offset_queue']
self.offset_confirmed_queue = offset_dict['offset_confirmed_queue']
except (FileNotFoundError, ValueError):
with open(constants.OFFSET_FILE, 'w') as f:
f.writelines(['offset_queue: 0',
'offset_confirmed_queue: 0'])
self.offset_queue = 0
self.offset_confirmed_queue = 0
def on_message(message: AMQPMessage, context: MessageContext, is_seed: bool):
def on_message(self, message: AMQPMessage, context: MessageContext, is_seed: bool):
annotation_message = AnnotationCreatedMessage(message)
print(f'Received: {annotation_message}')
save_folder = constants.data_seeds_dir if is_seed else constants.data_rabbit_dir
annotation_message.save_annotation(save_folder)
if is_seed:
self.offset_queue = context.offset
else:
self.offset_confirmed_queue = context.offset
with open(constants.OFFSET_FILE, 'w') as offset_file:
yaml.dump({'offset_queue': self.offset_queue,
'offset_confirmed_queue': self.offset_confirmed_queue}, offset_file)
await consumer.start()
await consumer.subscribe(stream=constants.ANNOTATIONS_QUEUE, callback=lambda msg, context: on_message(msg, context, is_seed=True),
offset_specification=ConsumerOffsetSpecification(OffsetType.OFFSET, offset_queue))
async def start(self):
await self.consumer.start()
await self.consumer.subscribe(stream=constants.ANNOTATIONS_QUEUE,
callback=lambda msg, context: self.on_message(msg, context, is_seed=True),
offset_specification=ConsumerOffsetSpecification(OffsetType.OFFSET, self.offset_queue))
await consumer.subscribe(stream=constants.ANNOTATIONS_CONFIRMED_QUEUE, callback=lambda msg, context: on_message(msg, context, is_seed=False),
offset_specification=ConsumerOffsetSpecification(OffsetType.OFFSET, offset_confirmed_queue))
await self.consumer.subscribe(stream=constants.ANNOTATIONS_CONFIRMED_QUEUE,
callback=lambda msg, context: self.on_message(msg, context, is_seed=False),
offset_specification=ConsumerOffsetSpecification(OffsetType.OFFSET, self.offset_confirmed_queue))
try:
await consumer.run()
await self.consumer.run()
except (KeyboardInterrupt, asyncio.CancelledError):
print("Closing Consumer...")
return
if __name__ == '__main__':
with asyncio.Runner() as runner:
runner.run(main())
runner.run(AnnotationQueueHandler().start())
dto/annotation_created_message.py
+49 -19
View File
@@ -1,6 +1,8 @@
import datetime
import os
import msgpack
import json
import datetime
class Detection:
@@ -17,21 +19,49 @@ class Detection:
return f'{self.cls}: {self.x:.2f} {self.y:.2f} {self.w:.2f} {self.h:.2f}, prob: {(self.confidence * 100):.1f}%'
class AnnotationCreatedMessage:
class AnnotationCreatedMessageNarrow:
def __init__(self, msgpack_bytes):
unpacked_data = msgpack.unpackb(msgpack_bytes)
self.createdDate = unpacked_data.get(0)
unpacked_data = msgpack.unpackb(msgpack_bytes, strict_map_key=False)
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.createdEmail = unpacked_data.get(2)
class AnnotationCreatedMessage:
last_offset = None
def __init__(self, msgpack_bytes):
unpacked_data = self.read_rabbit(msgpack_bytes)
ts = unpacked_data[0]
self.createdDate = datetime.datetime.utcfromtimestamp(ts.seconds) + datetime.timedelta(microseconds=ts.nanoseconds/1000)
self.name = unpacked_data[1]
self.originalMediaName = unpacked_data[2]
self.time = datetime.timedelta(microseconds=unpacked_data[3]/10)
self.imageExtension = unpacked_data[4]
detections_json_str = unpacked_data[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)
self.image = unpacked_data[6]
self.createdRole = unpacked_data[7]
self.createdEmail = unpacked_data[8]
self.source = unpacked_data[9]
self.status = unpacked_data[10]
@staticmethod
def read_rabbit(message_bytes):
if AnnotationCreatedMessage.last_offset is not None:
try:
unpacked_data = msgpack.unpackb(message_bytes[AnnotationCreatedMessage.last_offset:], raw=False, strict_map_key=False)
return unpacked_data
except Exception:
pass
for offset in range(3, 15):
try:
unpacked_data = msgpack.unpackb(message_bytes[offset:], raw=False, strict_map_key=False)
AnnotationCreatedMessage.last_offset = offset
return unpacked_data
except Exception:
pass
raise Exception(f'Cannot read rabbit message! Bytes: {message_bytes}')
def __str__(self):
if self.detections:
@@ -46,13 +76,13 @@ class AnnotationCreatedMessage:
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')
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 []
preprocessing.py
+4 -2
View File
@@ -28,6 +28,7 @@ transform = A.Compose([
A.HueSaturationValue(p=0.3, hue_shift_limit=8, sat_shift_limit=8, val_shift_limit=8)
], bbox_params=A.BboxParams(format='yolo'))
def correct_bboxes(labels):
margin = 0.0005
min_size = 0.01
@@ -125,10 +126,11 @@ def preprocess_annotations():
with concurrent.futures.ThreadPoolExecutor() as executor:
executor.map(process_image_file, images)
def process_image_file(image_file): # this function will be executed in thread
def process_image_file(image_file):
try:
image_path = os.path.join(data_images_dir, image_file.name)
labels_path = os.path.join(data_labels_dir, f'{Path(image_path).stem}.txt')
labels_path = os.path.join(data_labels_dir, f'{Path(str(image_path)).stem}.txt')
image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
img_ann = ImageLabel(