312 lines
11 KiB
Python
312 lines
11 KiB
Python
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)
|
|
|
|
|
|
app = FastAPI()
|
|
face_app = FastAPI()
|
|
app.mount("/face", face_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-face.pt" # 请替换为您的模型路径
|
|
KAFKA_BROKER = "222.186.10.253:9092"
|
|
KAFKA_TOPIC = "face" # 指定Kafka topic
|
|
KAFKA_GROUP_ID = "face_group" # 指定消费者组ID
|
|
|
|
REDIS_HOST = "222.186.10.253"
|
|
REDIS_PORT = 6379
|
|
REDIS_PASSWORD = "Obscura@2024"
|
|
REDIS_DB = 7
|
|
|
|
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 faceDetector:
|
|
def __init__(self, model_path):
|
|
self.model = YOLO(model_path).to('cuda:1')
|
|
|
|
def detect(self, frame):
|
|
results = self.model(frame, device='cuda:1')
|
|
return results
|
|
|
|
def format_results(self, results):
|
|
formatted_results = []
|
|
for r in results:
|
|
boxes = r.boxes
|
|
keypoints = r.keypoints
|
|
for i in range(len(boxes)):
|
|
box = boxes[i]
|
|
kpts = keypoints[i]
|
|
formatted_results.append({
|
|
"bbox": box.xyxy.tolist()[0],
|
|
"confidence": box.conf.item(),
|
|
"keypoints": kpts.xy.tolist()[0]
|
|
})
|
|
return formatted_results
|
|
|
|
def draw_results(self, frame, results, original_shape):
|
|
for r in results:
|
|
annotated_frame = r.plot(img=frame)
|
|
# 调整坐标以适应原始图像大小
|
|
h, w = annotated_frame.shape[:2]
|
|
scale_x, scale_y = original_shape[1] / w, original_shape[0] / h
|
|
annotated_frame = cv2.resize(annotated_frame, (original_shape[1], original_shape[0]))
|
|
return annotated_frame
|
|
|
|
detector = faceDetector(MODEL_PATH)
|
|
|
|
def process_image(image_data, filename):
|
|
try:
|
|
img = cv2.imdecode(np.frombuffer(image_data, np.uint8), cv2.IMREAD_COLOR)
|
|
original_shape = img.shape
|
|
# Convert BGR to RGB
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
|
|
# Resize image to fit model requirements (640x640)
|
|
img_resized = cv2.resize(img, (640, 640))
|
|
|
|
# Normalize and reshape to BCHW format
|
|
img_tensor = torch.from_numpy(img_resized).float().permute(2, 0, 1).unsqueeze(0) / 255.0
|
|
img_tensor = img_tensor.to('cuda:1')
|
|
|
|
results = detector.detect(img_tensor)
|
|
|
|
# Format results for JSON
|
|
json_results = detector.format_results(results)
|
|
|
|
# Draw results on original image
|
|
annotated_img = detector.draw_results(img_resized, results, original_shape)
|
|
|
|
# Save annotated image
|
|
annotated_filename = f"face_{filename}"
|
|
annotated_path = os.path.join(RESULT_DIR, annotated_filename)
|
|
cv2.imwrite(annotated_path, cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR))
|
|
|
|
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"face_{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))
|
|
original_shape = (height, width)
|
|
|
|
# Create output video file
|
|
annotated_filename = f"face_{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
|
|
|
|
# Process one frame per second
|
|
if frame_count % fps == 0:
|
|
# Convert BGR to RGB
|
|
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
|
|
# Resize frame to fit model requirements (640x640)
|
|
frame_resized = cv2.resize(frame_rgb, (640, 640))
|
|
|
|
# Normalize and reshape to BCHW format
|
|
frame_tensor = torch.from_numpy(frame_resized).float().permute(2, 0, 1).unsqueeze(0) / 255.0
|
|
frame_tensor = frame_tensor.to('cuda:1')
|
|
|
|
results = detector.detect(frame_tensor)
|
|
frame_json_results = detector.format_results(results)
|
|
json_results.append({"frame": frame_count, "detections": frame_json_results})
|
|
|
|
# Draw results on original frame
|
|
annotated_frame = detector.draw_results(frame_resized, results, original_shape)
|
|
# Convert RGB back to BGR for OpenCV
|
|
annotated_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)
|
|
else:
|
|
annotated_frame = frame
|
|
|
|
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
|
|
|
|
@face_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"face_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})
|
|
|
|
@face_app.get("/result/{filename}")
|
|
async def get_face_result(filename: str):
|
|
redis_key = f"face_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")
|
|
|
|
@face_app.get("/annotated/{filename}")
|
|
async def get_annotated_file(filename: str):
|
|
redis_key = f"face_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"face_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__:face_result:*') # 监听所有face_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"face_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=7004) |