import os import cv2 import torch import numpy as np from redis import Redis from ultralytics import YOLO import json from kafka import KafkaConsumer import threading import redis from config import * # 配置 MODEL_PATH = FALL_MODEL_PATH KAFKA_BROKER = KAFKA_BROKER KAFKA_TOPIC = WORKER_CONFIGS["fall"]["kafka_topic"] KAFKA_GROUP_ID = f"fall_{KAFKA_GROUP_ID_PREFIX}" REDIS_HOST = REDIS_HOST REDIS_PORT = REDIS_PORT REDIS_PASSWORD = REDIS_PASSWORD REDIS_DB = WORKER_CONFIGS["fall"]["redis_db"] # Worker使用的Redis DB MAIN_REDIS_DB = MAIN_REDIS_DB # 主Redis DB UPLOAD_DIR = UPLOAD_DIR RESULT_DIR = RESULT_DIR # 确保目录存在 os.makedirs(UPLOAD_DIR, exist_ok=True) os.makedirs(RESULT_DIR, exist_ok=True) # 初始化 Kafka consumer = KafkaConsumer( KAFKA_TOPIC, bootstrap_servers=[KAFKA_BROKER], group_id=KAFKA_GROUP_ID, auto_offset_reset='latest', enable_auto_commit=True, value_deserializer=lambda x: json.loads(x.decode('utf-8')) ) # 初始化 Redis redis_client = Redis( host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_DB ) main_redis_client = Redis( host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=MAIN_REDIS_DB ) class fallDetector: def __init__(self, model_path): self.model = YOLO(model_path) def detect(self, frame): results = self.model(frame) return results def format_results(self, results): formatted_results = [] for r in results: if not hasattr(r, 'boxes') or len(r.boxes) == 0: print("没有检测到任何对象") return [{"message": "No objects detected"}] boxes = r.boxes names = getattr(r, 'names', {}) for i in range(len(boxes)): box = boxes[i] if not hasattr(box, 'cls') or not hasattr(box, 'conf') or not hasattr(box, 'xyxy'): print(f"警告: 第 {i} 个框缺少必要的属性") continue try: class_id = int(box.cls.item()) formatted_result = { "bbox": box.xyxy.tolist()[0], "confidence": box.conf.item(), "class_id": class_id, "class": names.get(class_id, f"Unknown-{class_id}") } formatted_results.append(formatted_result) except Exception as e: print(f"处理第 {i} 个框时出错: {str(e)}") # print("格式化后的结果:", formatted_results) return formatted_results def draw_results(self, frame, results): for r in results: annotated_frame = r.plot() return annotated_frame detector = fallDetector(MODEL_PATH) def process_image(image_data, filename): try: img = cv2.imdecode(np.frombuffer(image_data, np.uint8), cv2.IMREAD_COLOR) results = detector.detect(img) # Format results for JSON json_results = detector.format_results(results) # Draw results on image annotated_img = detector.draw_results(img, results) # Save annotated image annotated_filename = f"fall_{filename}" annotated_path = os.path.join(RESULT_DIR, annotated_filename) cv2.imwrite(annotated_path, annotated_img) return json_results, annotated_filename except Exception as e: print(f"Error processing image: {str(e)}") return None, None def process_video(video_data, filename): try: temp_video_path = os.path.join(UPLOAD_DIR, f"fall_{filename}") with open(temp_video_path, 'wb') as temp_video: temp_video.write(video_data) cap = cv2.VideoCapture(temp_video_path) frame_count = 0 json_results = [] # Get video properties fps = int(cap.get(cv2.CAP_PROP_FPS)) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) # Create output video file annotated_filename = f"fall_{filename}" output_path = os.path.join(RESULT_DIR, annotated_filename) out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height)) while cap.isOpened(): ret, frame = cap.read() if not ret: break results = detector.detect(frame) frame_json_results = detector.format_results(results) json_results.append({"frame": frame_count, "detections": frame_json_results}) annotated_frame = detector.draw_results(frame, results) out.write(annotated_frame) frame_count += 1 cap.release() out.release() # Clean up temporary input video file os.remove(temp_video_path) return json_results, annotated_filename except Exception as e: print(f"Error processing video: {str(e)}") return None, None def process_task(): print("开始处理任务,等待Kafka消息...") for message in consumer: print(f"收到Kafka消息: topic={message.topic}, partition={message.partition}, offset={message.offset}") task = message.value task_id = task['task_id'] filename = task['filename'] file_type = task['file_type'] print(f"解析任务信息: ID={task_id}, 文件名={filename}, 类型={file_type}") file_path = os.path.join(UPLOAD_DIR, filename) # 检查键类型并更新状态 task_key = f"task:{task_id}" try: key_type = main_redis_client.type(task_key) if key_type != b'hash': main_redis_client.delete(task_key) main_redis_client.hset(f"task:{task_id}", "status", "processing") print(f"任务 {task_id} 状态更新为 'processing'") except redis.exceptions.ResponseError as e: print(f"更新任务 {task_id} 状态时出错: {str(e)}") continue # 跳过这个任务，继续处理下一个 try: if file_type == "image": print(f"开始处理图像: {filename}") with open(file_path, 'rb') as f: image_data = f.read() json_results, annotated_filename = process_image(image_data, filename) if json_results and annotated_filename is not None: result_path = os.path.join(RESULT_DIR, annotated_filename) redis_client.hset(f"fall_result:{task_id}", mapping={ "result": json.dumps(json_results), "result_file": annotated_filename }) main_redis_client.hset(f"task:{task_id}", "status", "completed") main_redis_client.hset(f"task:{task_id}", "result_type", "fall") main_redis_client.hset(f"task:{task_id}", "result_key", f"fall_result:{task_id}") print(f"图像 {filename} 处理完成,结果已保存") else: print(f"图像 {filename} 处理失败") main_redis_client.hset(f"task:{task_id}", "status", "failed") else: # video print(f"开始处理视频: {filename}") with open(file_path, 'rb') as f: video_data = f.read() json_results, annotated_filename = process_video(video_data, filename) if json_results and annotated_filename: redis_client.hset(f"fall_result:{task_id}", mapping={ "result": json.dumps(json_results), "result_file": annotated_filename }) main_redis_client.hset(f"task:{task_id}", "status", "completed") main_redis_client.hset(f"task:{task_id}", "result_type", "fall") main_redis_client.hset(f"task:{task_id}", "result_key", f"fall_result:{task_id}") print(f"视频 {filename} 处理完成,结果已保存") else: print(f"视频 {filename} 处理失败") main_redis_client.hset(f"task:{task_id}", "status", "failed") except Exception as e: print(f"处理任务 {task_id} 时出错: {str(e)}") main_redis_client.hset(f"task:{task_id}", { "status": "failed", "error": str(e) }) print(f"任务 {task_id} 处理完毕,等待下一个Kafka消息...") def listen_redis_changes(): pubsub = redis_client.pubsub() pubsub.psubscribe('__keyspace@3__:fall_result:*') # 监听所有fall_result键的变化 for message in pubsub.listen(): if message['type'] == 'pmessage': key = message['channel'].decode('utf-8').split(':')[-1] operation = message['data'].decode('utf-8') if operation == 'hset': value = redis_client.hgetall(f"fall_result:{key}") if value: result = {k.decode(): v.decode() for k, v in value.items()} print(f"Result update for task {key}: {result}") if __name__ == "__main__": print("fall处理程序启动...") # 启动处理任务的线程 task_thread = threading.Thread(target=process_task, daemon=True) task_thread.start() print("任务处理线程已启动") # 启动Redis监听线程 redis_thread = threading.Thread(target=listen_redis_changes, daemon=True) redis_thread.start() print("Redis监听线程已启动") print("主程序进入等待状态...") # 保持主线程运行 task_thread.join() redis_thread.join()