from fastapi import FastAPI, HTTPException, File, UploadFile from fastapi.responses import StreamingResponse, JSONResponse from fastapi.middleware.cors import CORSMiddleware from ultralytics import YOLO import cv2 import numpy as np import json import uvicorn from kafka import KafkaProducer, KafkaConsumer from redis import Redis import io import uuid import os from datetime import datetime, timedelta import threading app = FastAPI() fall_app = FastAPI() app.mount("/fall", fall_app) # CORS配置 ALLOWED_ORIGINS = 'https://beta.obscura.work' # 只为主应用添加CORS中间件 app.add_middleware( CORSMiddleware, allow_origins=ALLOWED_ORIGINS, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # 配置 MODEL_PATH = "/home/zydi/models/yolov8n-fall.pt" # 请替换为您的模型路径 KAFKA_BROKER = "222.186.10.253:9092" KAFKA_TOPIC = "fall" # 指定Kafka topic KAFKA_GROUP_ID = "fall_group" # 指定消费者组ID REDIS_HOST = "222.186.10.253" REDIS_PORT = 6379 REDIS_PASSWORD = "Obscura@2024" REDIS_DB = 4 UPLOAD_DIR = "/www/wwwroot/beta.obscura.work/upload_files/upload" RESULT_DIR = "/www/wwwroot/beta.obscura.work/upload_files/result" MAX_FILE_AGE = timedelta(hours=1) # 确保目录存在 os.makedirs(UPLOAD_DIR, exist_ok=True) os.makedirs(RESULT_DIR, exist_ok=True) # 初始化 Kafka producer = KafkaProducer(bootstrap_servers=[KAFKA_BROKER]) consumer = KafkaConsumer( KAFKA_TOPIC, bootstrap_servers=[KAFKA_BROKER], group_id=KAFKA_GROUP_ID, auto_offset_reset='earliest', 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 ) 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 @fall_app.post("/upload") async def upload_file(file: UploadFile = File(...)): content = await file.read() file_extension = os.path.splitext(file.filename)[1].lower() new_filename = f"{uuid.uuid4()}{file_extension}" # Save the original file original_file_path = os.path.join(UPLOAD_DIR, new_filename) with open(original_file_path, "wb") as f: f.write(content) # Send processing task to Kafka producer.send(KAFKA_TOPIC, json.dumps({ "filename": new_filename, "file_type": "image" if file_extension in ['.jpg', '.jpeg', '.png'] else "video" }).encode('utf-8')) # Set initial status in Redis redis_key = f"fall_result:{new_filename}" redis_client.set(redis_key, json.dumps({"status": "queued"})) return JSONResponse(content={"message": "File uploaded and queued for processing", "filename": new_filename}) @fall_app.get("/result/{filename}") async def get_fall_result(filename: str): redis_key = f"fall_result:{filename}" result = redis_client.get(redis_key) if result: result_data = json.loads(result) return JSONResponse(content=result_data) # 直接返回整个结果，包括 status else: raise HTTPException(status_code=404, detail="Result not found") @fall_app.get("/annotated/{filename}") async def get_annotated_file(filename: str): redis_key = f"fall_result:{filename}" result = redis_client.get(redis_key) if result: result_data = json.loads(result) if result_data["status"] == "completed": annotated_filename = result_data["annotated_filename"] file_path = os.path.join(RESULT_DIR, annotated_filename) if os.path.exists(file_path): def iterfile(): with open(file_path, mode="rb") as file_like: yield from file_like file_extension = os.path.splitext(annotated_filename)[1].lower() return StreamingResponse(iterfile(), media_type=f"image/{file_extension[1:]}" if file_extension in ['.jpg', '.jpeg', '.png'] else "video/mp4") raise HTTPException(status_code=404, detail="Annotated file not found") def process_task(): for message in consumer: task = message.value filename = task['filename'] file_type = task['file_type'] file_path = os.path.join(UPLOAD_DIR, filename) # Update status to "processing" redis_key = f"fall_result:{filename}" redis_client.set(redis_key, json.dumps({"status": "processing"})) try: if file_type == "image": with open(file_path, 'rb') as f: content = f.read() json_results, annotated_filename = process_image(content, filename) else: with open(file_path, 'rb') as f: content = f.read() json_results, annotated_filename = process_video(content, filename) if json_results and annotated_filename: redis_client.set(redis_key, json.dumps({ "json_results": json_results, "status": "completed", "annotated_filename": annotated_filename })) else: redis_client.set(redis_key, json.dumps({"status": "failed"})) except Exception as e: print(f"Error processing task: {str(e)}") redis_client.set(redis_key, json.dumps({"status": "failed", "error": str(e)})) 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 == 'set': value = redis_client.get(f"fall_result:{key}") if value: result = json.loads(value) print(f"Status update for {key}: {result['status']}") # 这里可以添加其他处理逻辑，比如发送通知等 if __name__ == "__main__": # 启动处理任务的线程 threading.Thread(target=process_task, daemon=True).start() # 启动Redis监听线程 threading.Thread(target=listen_redis_changes, daemon=True).start() uvicorn.run(app, host="0.0.0.0", port=7002)