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 import torch torch.cuda.set_device(1) import colorsys app = FastAPI() yolo_app = FastAPI() app.mount("/yolo", yolo_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/yolov8x.pt" # 请替换为您的模型路径 KAFKA_BROKER = "222.186.10.253:9092" KAFKA_TOPIC = "yolo" # 指定Kafka topic KAFKA_GROUP_ID = "yolo_group" # 指定消费者组ID REDIS_HOST = "222.186.10.253" REDIS_PORT = 6379 REDIS_PASSWORD = "Obscura@2024" REDIS_DB = 6 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 yoloDetector: 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, original_shape): formatted_results = [] for r in results: boxes = r.boxes for box in boxes: x1, y1, x2, y2 = box.xyxy[0].tolist() # 缩放坐标到原始图像尺寸 x1, x2 = [x * original_shape[1] / 640 for x in [x1, x2]] y1, y2 = [y * original_shape[0] / 640 for y in [y1, y2]] conf = box.conf.item() cls = int(box.cls.item()) name = self.model.names[cls] formatted_results.append({ "class": name, "confidence": conf, "bbox": [x1, y1, x2, y2] }) return formatted_results def draw_results(self, frame, formatted_results): for result in formatted_results: x1, y1, x2, y2 = map(int, result['bbox']) name = result['class'] conf = result['confidence'] cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) # 使用固定的绿色 label = f"{name} {conf:.2f}" (text_width, text_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 1, 2) cv2.rectangle(frame, (x1, y1 - text_height - 5), (x1 + text_width, y1), (0, 255, 0), -1) cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2) return frame detector = yoloDetector(MODEL_PATH) def process_image(image_data, filename): try: original_img = cv2.imdecode(np.frombuffer(image_data, np.uint8), cv2.IMREAD_COLOR) original_shape = original_img.shape img = cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB) img = cv2.resize(img, (640, 640)) img = img.transpose((2, 0, 1)) img = np.ascontiguousarray(img) img = torch.from_numpy(img).float() img /= 255.0 img = img.unsqueeze(0) results = detector.detect(img) json_results = detector.format_results(results, original_shape) annotated_img = detector.draw_results(original_img, json_results) annotated_filename = f"yolo_{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"处理图像时出错: {str(e)}") return None, None def process_video(video_data, filename): try: temp_video_path = os.path.join(UPLOAD_DIR, f"yolo_{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 = [] 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)) original_shape = (height, width) annotated_filename = f"yolo_{filename}" output_path = os.path.join(RESULT_DIR, annotated_filename) out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 1, (width, height)) while cap.isOpened(): ret, frame = cap.read() if not ret: break # 每秒只处理一帧 if frame_count % fps == 0: preprocessed_frame = preprocess_frame(frame) results = detector.detect(preprocessed_frame) frame_json_results = detector.format_results(results, original_shape) json_results.append({"frame": frame_count, "detections": frame_json_results}) annotated_frame = detector.draw_results(frame, frame_json_results) out.write(annotated_frame) frame_count += 1 cap.release() out.release() os.remove(temp_video_path) return json_results, annotated_filename except Exception as e: print(f"处理视频时出错: {str(e)}") return None, None def preprocess_frame(frame): # 预处理单个视频帧 frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame_resized = cv2.resize(frame_rgb, (640, 640)) # 调整为YOLO输入尺寸 frame_transposed = frame_resized.transpose((2, 0, 1)) # HWC转为CHW frame_contiguous = np.ascontiguousarray(frame_transposed) frame_tensor = torch.from_numpy(frame_contiguous).float() frame_normalized = frame_tensor / 255.0 # 归一化到[0, 1] frame_batched = frame_normalized.unsqueeze(0) # 添加批次维度 return frame_batched @yolo_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"yolo_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}) @yolo_app.get("/result/{filename}") async def get_yolo_result(filename: str): redis_key = f"yolo_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") @yolo_app.get("/annotated/{filename}") async def get_annotated_file(filename: str): redis_key = f"yolo_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"yolo_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__:yolo_result:*') # 监听所有yolo_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"yolo_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=7003)