318 lines
12 KiB
Python
318 lines
12 KiB
Python
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
|
|
import torch
|
|
from config import *
|
|
import insightface
|
|
from insightface.app import FaceAnalysis
|
|
from insightface.utils import face_align
|
|
|
|
|
|
torch.cuda.set_device(1)
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
# 配置
|
|
|
|
KAFKA_BROKER = KAFKA_BROKER
|
|
KAFKA_TOPIC = WORKER_CONFIGS["compare"]["kafka_topic"]
|
|
KAFKA_GROUP_ID = f"compare_{KAFKA_GROUP_ID_PREFIX}"
|
|
|
|
REDIS_HOST = REDIS_HOST
|
|
REDIS_PORT = REDIS_PORT
|
|
REDIS_PASSWORD = REDIS_PASSWORD
|
|
REDIS_DB = WORKER_CONFIGS["compare"]["redis_db"] # Worker使用的Redis DB
|
|
MAIN_REDIS_DB = MAIN_REDIS_DB # 主Redis DB
|
|
|
|
UPLOAD_DIR = UPLOAD_DIR
|
|
RESULT_DIR = RESULT_DIR
|
|
|
|
# 初始化 Kafka
|
|
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
|
|
)
|
|
|
|
main_redis_client = Redis(
|
|
host=REDIS_HOST,
|
|
port=REDIS_PORT,
|
|
password=REDIS_PASSWORD,
|
|
db=MAIN_REDIS_DB
|
|
)
|
|
|
|
|
|
class FaceComparator:
|
|
def __init__(self):
|
|
print("初始化 InsightFace...")
|
|
self.app = FaceAnalysis(name='buffalo_l', allowed_modules=['detection', 'recognition', 'genderage'])
|
|
self.app.prepare(ctx_id=0, det_size=(640, 640))
|
|
print(f"InsightFace 初始化完成，使用设备: {device}")
|
|
|
|
|
|
def detect(self, frame):
|
|
"""检测人脸并返回所有特征"""
|
|
try:
|
|
faces = self.app.get(frame)
|
|
return faces
|
|
except Exception as e:
|
|
return []
|
|
|
|
def format_results(self, faces):
|
|
"""格式化检测结果"""
|
|
print("开始格式化检测结果...")
|
|
try:
|
|
formatted_results = []
|
|
for i, face in enumerate(faces):
|
|
# 设置默认值，避免None导致的错误
|
|
result = {
|
|
'bbox': face.bbox.tolist(),
|
|
'kps': face.kps.tolist(),
|
|
'gender': int(face.gender) if hasattr(face, 'gender') and face.gender is not None else -1,
|
|
'age': float(face.age) if hasattr(face, 'age') and face.age is not None else 0.0,
|
|
'det_score': float(face.det_score),
|
|
'embedding': face.embedding.tolist()
|
|
}
|
|
formatted_results.append(result)
|
|
|
|
return formatted_results
|
|
except Exception as e:
|
|
print(f"格式化结果时出错: {str(e)}")
|
|
print(f"错误类型: {type(e)}")
|
|
import traceback
|
|
print(f"错误堆栈: {traceback.format_exc()}")
|
|
return []
|
|
|
|
def draw_results(self, frame, results):
|
|
"""在图像上绘制检测结果"""
|
|
print("开始绘制检测结果...")
|
|
try:
|
|
annotated_frame = frame.copy()
|
|
|
|
if not results:
|
|
print("没有检测结果，返回原始图像")
|
|
return annotated_frame
|
|
|
|
for i, r in enumerate(results):
|
|
bbox = r['bbox']
|
|
kps = r['kps']
|
|
|
|
cv2.rectangle(annotated_frame,
|
|
(int(bbox[0]), int(bbox[1])),
|
|
(int(bbox[2]), int(bbox[3])),
|
|
(0, 255, 0), 2)
|
|
|
|
for x, y in kps:
|
|
cv2.circle(annotated_frame,
|
|
(int(x), int(y)),
|
|
3, (255, 255, 0), -1)
|
|
|
|
gender_text = 'Male' if r['gender'] == 1 else 'Female' if r['gender'] == 0 else 'Unknown'
|
|
label = f"Age: {int(r['age'])} Gender: {gender_text}"
|
|
cv2.putText(annotated_frame,
|
|
label,
|
|
(int(bbox[0]), int(bbox[1]-10)),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.5,
|
|
(0, 255, 0),
|
|
2)
|
|
return annotated_frame
|
|
|
|
except Exception as e:
|
|
print(f"绘制结果时出错: {str(e)}")
|
|
return frame.copy()
|
|
|
|
detector = FaceComparator()
|
|
|
|
def process_image(image_path):
|
|
try:
|
|
original_img = cv2.imread(image_path)
|
|
if original_img is None:
|
|
raise ValueError(f"无法读取图像文件: {image_path}")
|
|
|
|
if original_img.size == 0:
|
|
raise ValueError("图像数据为空")
|
|
|
|
faces = detector.detect(original_img)
|
|
|
|
if not faces:
|
|
json_results = []
|
|
else:
|
|
json_results = detector.format_results(faces)
|
|
if not json_results: # 如果格式化失败，使用空列表
|
|
json_results = []
|
|
|
|
annotated_img = detector.draw_results(original_img, json_results)
|
|
|
|
# 即使没有检测到人脸或格式化失败，也返回处理结果
|
|
return json_results, annotated_img
|
|
|
|
except Exception as e:
|
|
return None, None
|
|
|
|
def process_video(video_path):
|
|
try:
|
|
cap = cv2.VideoCapture(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)
|
|
|
|
out = cv2.VideoWriter(video_path.replace(UPLOAD_DIR, RESULT_DIR), cv2.VideoWriter_fourcc(*'mp4v'), fps, (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()
|
|
|
|
return json_results
|
|
except Exception as e:
|
|
print(f"处理视频时出错: {str(e)}")
|
|
return None
|
|
|
|
def preprocess_frame(frame):
|
|
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
frame_resized = cv2.resize(frame_rgb, (640, 640))
|
|
frame_transposed = frame_resized.transpose((2, 0, 1))
|
|
frame_contiguous = np.ascontiguousarray(frame_transposed)
|
|
frame_tensor = torch.from_numpy(frame_contiguous).float()
|
|
frame_normalized = frame_tensor.to(device) / 255.0
|
|
frame_batched = frame_normalized.unsqueeze(0)
|
|
return frame_batched
|
|
|
|
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(task_key, "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}")
|
|
json_results, annotated_img = process_image(file_path)
|
|
if json_results and annotated_img is not None:
|
|
result_filename = f"compare_{filename}"
|
|
result_path = os.path.join(RESULT_DIR, result_filename)
|
|
cv2.imwrite(result_path, annotated_img)
|
|
|
|
redis_client.hset(f"compare_result:{task_id}",
|
|
"result", json.dumps(json_results))
|
|
redis_client.hset(f"compare_result:{task_id}",
|
|
"result_file", result_filename)
|
|
|
|
main_redis_client.hset(f"task:{task_id}", "status", "completed")
|
|
main_redis_client.hset(f"task:{task_id}", "result_type", "compare")
|
|
main_redis_client.hset(f"task:{task_id}", "result_key", f"compare_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}")
|
|
json_results = process_video(file_path)
|
|
if json_results:
|
|
result_filename = f"compare_{filename}"
|
|
redis_client.hset(f"compare_result:{task_id}",
|
|
"result", json.dumps(json_results))
|
|
redis_client.hset(f"compare_result:{task_id}",
|
|
"result_file", result_filename)
|
|
|
|
main_redis_client.hset(f"task:{task_id}", "status", "completed")
|
|
main_redis_client.hset(f"task:{task_id}", "result_type", "compare")
|
|
main_redis_client.hset(f"task:{task_id}", "result_key", f"compare_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")
|
|
main_redis_client.hset(f"task:{task_id}", "error", str(e))
|
|
|
|
print(f"任务 {task_id} 处理完毕,等待下一个Kafka消息...")
|
|
def listen_redis_changes():
|
|
pubsub = redis_client.pubsub()
|
|
pubsub.psubscribe('__keyspace@3__:compare_result:*') # 监听所有compare_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 == 'hmset':
|
|
value = redis_client.hgetall(f"compare_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("compare处理程序启动...")
|
|
# 启动处理任务的线程
|
|
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() |