300 lines
13 KiB
Python
300 lines
13 KiB
Python
|
import ctypes
|
|||
|
|
import datetime
|
|||
|
|
import json
|
|||
|
|
import os
|
|||
|
|
import time
|
|||
|
|
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
|
|||
|
|
from cv_back_video import extract_box_details, cal_tricker_results
|
|||
|
|
from middleware.minio_util import upload_file, downFile
|
|||
|
|
from middleware.read_pic_metadata import read_dji_exif_to_dict
|
|||
|
|
from mqtt_pub import MQTTClient
|
|||
|
|
from ultralytics import solutions
|
|||
|
|
|
|||
|
|
from yolo.detect.yolo11_det_pic_trt import YoLo11TRT
|
|||
|
|
|
|||
|
|
mqtt_client = None
|
|||
|
|
# MQTT 代理地址和端口
|
|||
|
|
# broker = "112.44.103.230" # 公共 MQTT 代理（免费）
|
|||
|
|
broker = "8.137.54.85" # 公共 MQTT 代理（免费）
|
|||
|
|
port = 1883 # MQTT 默认端口
|
|||
|
|
# 主题
|
|||
|
|
topic = "thing/product/ai/events"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def pic_detect_func(mqtt_pub_ip, mqtt_pub_port, mqtt_pub_topic,
|
|||
|
|
task_id, model_func_id_list, para, s3_url_list, m1, cls, chinese_label, confidence,
|
|||
|
|
use_fp16=False):
|
|||
|
|
global mqtt_client
|
|||
|
|
ov_model = None
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 设置环境变量，提高YOLO性能
|
|||
|
|
os.environ["OMP_NUM_THREADS"] = "4"
|
|||
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
|||
|
|
|
|||
|
|
# 导入必要的库
|
|||
|
|
import torch
|
|||
|
|
print(f"PyTorch 可用: {torch.__version__}, CUDA可用: {torch.cuda.is_available()}")
|
|||
|
|
if torch.cuda.is_available():
|
|||
|
|
print(f"GPU: {torch.cuda.get_device_name(0)}")
|
|||
|
|
torch.cuda.empty_cache()
|
|||
|
|
|
|||
|
|
# 预加载YOLO模型
|
|||
|
|
print("预加载YOLO模型...")
|
|||
|
|
max_retries = 3
|
|||
|
|
retry_count = 0
|
|||
|
|
model_params = {}
|
|||
|
|
# 尝试检测YOLO版本以适配不同版本的参数
|
|||
|
|
try:
|
|||
|
|
# 先简单尝试加载模型，不带任何参数
|
|||
|
|
test_model = YOLO(m1)
|
|||
|
|
# 如果成功，检查可用的参数
|
|||
|
|
if hasattr(test_model, "task"):
|
|||
|
|
model_params["task"] = "detect" # 指定任务类型
|
|||
|
|
|
|||
|
|
# 检查是否支持half精度
|
|||
|
|
if torch.cuda.is_available():
|
|||
|
|
model_params["half"] = True
|
|||
|
|
|
|||
|
|
# 检查是否支持verbose参数
|
|||
|
|
import inspect
|
|||
|
|
if "verbose" in inspect.signature(YOLO.__init__).parameters:
|
|||
|
|
model_params["verbose"] = False
|
|||
|
|
|
|||
|
|
print(f"检测到支持的YOLO参数: {model_params}")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"参数检测失败，将使用默认参数: {e}")
|
|||
|
|
model_params = {}
|
|||
|
|
|
|||
|
|
while retry_count < max_retries:
|
|||
|
|
try:
|
|||
|
|
ov_model = YOLO(m1, **model_params)
|
|||
|
|
dummy_frame = np.zeros((1080, 1920, 3), dtype=np.uint8)
|
|||
|
|
for _ in range(3):
|
|||
|
|
ov_model(dummy_frame, classes=cls, conf=confidence, show=False)
|
|||
|
|
print("YOLO模型加载成功并预热完成")
|
|||
|
|
break
|
|||
|
|
except Exception as e:
|
|||
|
|
retry_count += 1
|
|||
|
|
print(f"YOLO模型加载失败 (尝试 {retry_count}/{max_retries}): {e}")
|
|||
|
|
if "got an unexpected keyword argument" in str(e) and model_params:
|
|||
|
|
param_name = str(e).split("'")[-2]
|
|||
|
|
if param_name in model_params:
|
|||
|
|
del model_params[param_name]
|
|||
|
|
time.sleep(2)
|
|||
|
|
|
|||
|
|
if ov_model is None:
|
|||
|
|
raise Exception("无法加载YOLO模型，达到最大重试次数")
|
|||
|
|
|
|||
|
|
# 优化模型推理参数
|
|||
|
|
ov_model.conf = confidence
|
|||
|
|
|
|||
|
|
# 将模型移到设备（GPU或CPU）
|
|||
|
|
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
|
|||
|
|
try:
|
|||
|
|
ov_model.to(device)
|
|||
|
|
# 调整批处理大小为1，减少内存占用
|
|||
|
|
if hasattr(ov_model, 'args') and hasattr(ov_model.args, 'batch'):
|
|||
|
|
ov_model.args.batch = 1
|
|||
|
|
# 启用半精度
|
|||
|
|
if use_fp16 and device.startswith('cuda') and hasattr(ov_model, 'model'):
|
|||
|
|
try:
|
|||
|
|
ov_model.model = ov_model.model.half()
|
|||
|
|
print("✅ 启用半精度 FP16 模式")
|
|||
|
|
except Exception as half_err:
|
|||
|
|
print(f"⚠️ 半精度转换失败: {half_err}")
|
|||
|
|
else:
|
|||
|
|
print("ℹ️ 半精度模式未启用")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"⚠️ 模型设备配置警告: {e}")
|
|||
|
|
|
|||
|
|
mqtt_client = MQTTClient(mqtt_pub_ip, mqtt_pub_port, mqtt_pub_topic)
|
|||
|
|
|
|||
|
|
# solutions默认会在图上绘制一个方框，很烦
|
|||
|
|
line_points = [(-10, -10), (-10, -10)]
|
|||
|
|
# 伪代码，后续识别类型需要跟ai_model_list 的id关联
|
|||
|
|
# cls = [0, 1, 2, 3, 4]
|
|||
|
|
# 这里参数show=True，一旦通过rest接口重启程序，会导致线程锁住
|
|||
|
|
|
|||
|
|
counter = solutions.ObjectCounter(show=False, region=line_points, model=ov_model, classes=cls)
|
|||
|
|
|
|||
|
|
# Model processing setup
|
|||
|
|
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
|
|||
|
|
try:
|
|||
|
|
counter.model.to(device)
|
|||
|
|
if hasattr(counter.model, 'args') and hasattr(counter.model.args, 'batch'):
|
|||
|
|
counter.model.args.batch = 1
|
|||
|
|
# 启用半精度
|
|||
|
|
if use_fp16 and device.startswith('cuda') and hasattr(counter.model, 'model'):
|
|||
|
|
try:
|
|||
|
|
counter.model.model = counter.model.model.half()
|
|||
|
|
print("✅ 启用半精度 FP16 模式")
|
|||
|
|
except Exception as half_err:
|
|||
|
|
print(f"⚠️ 半精度转换失败: {half_err}")
|
|||
|
|
else:
|
|||
|
|
print("ℹ️ 半精度模式未启用")
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"⚠️ 模型设备配置警告: {e}")
|
|||
|
|
error_count = 0
|
|||
|
|
|
|||
|
|
# Get class names from the model
|
|||
|
|
class_names = counter.model.names if hasattr(counter.model, 'names') else {i: str(i) for i in range(1000)}
|
|||
|
|
|
|||
|
|
local_func_cache = {
|
|||
|
|
"func_100000": None,
|
|||
|
|
"func_100004": None, # 存储缓存，缓存人员track_id
|
|||
|
|
"func_100006": None # 存储缓存，缓存车辆track_id
|
|||
|
|
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
for pic_url in s3_url_list:
|
|||
|
|
image_path = downFile(pic_url)
|
|||
|
|
print("1")
|
|||
|
|
frame = cv2.imread(image_path)
|
|||
|
|
|
|||
|
|
# 读取照片元信息，包括经纬度、飞行姿态
|
|||
|
|
image_exif_dict = read_dji_exif_to_dict(image_path)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
frame_copy = frame.copy()
|
|||
|
|
results = counter(frame)
|
|||
|
|
annotated_frame, box_result = cal_tricker_results(frame_copy, counter, class_names,
|
|||
|
|
model_func_id_list,
|
|||
|
|
local_func_cache, para, cls, chinese_label,
|
|||
|
|
model_func_id_list[0])
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"处理帧错误: {e}")
|
|||
|
|
error_count += 1
|
|||
|
|
if error_count >= 5:
|
|||
|
|
print(f"连续处理错误达到5次 ，正在停止处理...")
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
# 提取边界框统计信息
|
|||
|
|
# box_detail = extract_box_details(last_results, model_id)
|
|||
|
|
# 保存结果图片
|
|||
|
|
local_path = f"output_frames/{os.path.basename(image_path)}"
|
|||
|
|
success = cv2.imwrite(local_path, annotated_frame, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
|
|||
|
|
if success:
|
|||
|
|
print(f"图像已保存到: {local_path}")
|
|||
|
|
else:
|
|||
|
|
print("错误: 图像保存失败")
|
|||
|
|
date_str = str(datetime.datetime.now().timestamp())
|
|||
|
|
|
|||
|
|
minio_path_before, file_type_before = upload_file(image_path, None)
|
|||
|
|
minio_path, file_type = upload_file(local_path, date_str)
|
|||
|
|
os.remove(local_path)
|
|||
|
|
message = {
|
|||
|
|
"task_id": task_id,
|
|||
|
|
"minio": {
|
|||
|
|
"minio_path": minio_path,
|
|||
|
|
"file_type": file_type,
|
|||
|
|
"original_path": minio_path_before
|
|||
|
|
},
|
|||
|
|
"box_detail": box_result,
|
|||
|
|
"image_exif_dict": image_exif_dict
|
|||
|
|
}
|
|||
|
|
# print("16")
|
|||
|
|
# json_message = json.dumps(message)
|
|||
|
|
# 将字典转换为 JSON 字符串，使用双引号，并保留中文字符
|
|||
|
|
json_message = json.dumps(message, indent=4, ensure_ascii=False)
|
|||
|
|
mqtt_client.publish_message(json_message)
|
|||
|
|
time.sleep(0.1)
|
|||
|
|
if os.path.exists(minio_path):
|
|||
|
|
os.remove(minio_path)
|
|||
|
|
if os.path.exists(minio_path_before):
|
|||
|
|
os.remove(minio_path_before)
|
|||
|
|
except Exception as infer_err:
|
|||
|
|
print(f"推理错误: {infer_err}")
|
|||
|
|
except Exception as infer_err:
|
|||
|
|
print(f"pic_detect 推理错误: {infer_err}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def pic_detect_func_trt(mqtt_pub_ip, mqtt_pub_port, mqtt_pub_topic,
|
|||
|
|
task_id, s3_url_list, engine_file_path,
|
|||
|
|
PLUGIN_LIBRARY, confidence, categories ):
|
|||
|
|
global mqtt_client
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
print(f"mqtt_client {mqtt_pub_ip} {mqtt_pub_port} {mqtt_pub_topic}")
|
|||
|
|
mqtt_client = MQTTClient(mqtt_pub_ip, mqtt_pub_port, mqtt_pub_topic)
|
|||
|
|
print("mqtt")
|
|||
|
|
# # 设置引擎和插件路径
|
|||
|
|
# engine_file_path = r"D:\project\AI-PYTHON\Ai_tottle\engine\gdaq_hat_0926\gdaq_hat_0926.engine"
|
|||
|
|
# PLUGIN_LIBRARY = r"D:\project\AI-PYTHON\Ai_tottle\engine\gdaq_hat_0926\myplugins.dll"
|
|||
|
|
|
|||
|
|
# 加载插件
|
|||
|
|
ctypes.CDLL(PLUGIN_LIBRARY)
|
|||
|
|
print(f"PLUGIN_LIBRARY{PLUGIN_LIBRARY}")
|
|||
|
|
# 初始化推理器
|
|||
|
|
yolo11_wrapper = YoLo11TRT(engine_file_path)
|
|||
|
|
print(f"engine_file_path{engine_file_path} {len(s3_url_list)}")
|
|||
|
|
try:
|
|||
|
|
for pic_url in s3_url_list:
|
|||
|
|
image_path = downFile(pic_url)
|
|||
|
|
print("1")
|
|||
|
|
frame = cv2.imread(image_path)
|
|||
|
|
print("2")
|
|||
|
|
# # 读取照片元信息，包括经纬度、飞行姿态
|
|||
|
|
# image_exif_dict = read_dji_exif_to_dict(image_path)
|
|||
|
|
batch_image_raw, result_box_list, result_scores_list, result_classid_list, infer_time = yolo11_wrapper.infer(
|
|||
|
|
[frame], confidence, categories)
|
|||
|
|
detected_frame = batch_image_raw[0]
|
|||
|
|
|
|||
|
|
# 提取边界框统计信息
|
|||
|
|
# box_detail = extract_box_details(last_results, model_id)
|
|||
|
|
# 保存结果图片
|
|||
|
|
local_path = f"output_frames/{os.path.basename(image_path)}"
|
|||
|
|
success = cv2.imwrite(local_path, detected_frame, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
|
|||
|
|
if success:
|
|||
|
|
print(f"图像已保存到: {local_path}")
|
|||
|
|
else:
|
|||
|
|
print("错误: 图像保存失败")
|
|||
|
|
date_str = str(datetime.datetime.now().timestamp())
|
|||
|
|
|
|||
|
|
minio_path_before, file_type_before = upload_file(image_path, None)
|
|||
|
|
minio_path, file_type = upload_file(local_path, date_str)
|
|||
|
|
os.remove(local_path)
|
|||
|
|
message = {
|
|||
|
|
"task_id": task_id,
|
|||
|
|
"minio": {
|
|||
|
|
"minio_path": minio_path,
|
|||
|
|
"file_type": file_type,
|
|||
|
|
"original_path": minio_path_before
|
|||
|
|
},
|
|||
|
|
"box_detail": {
|
|||
|
|
"result_box_list": result_box_list,
|
|||
|
|
"result_scores_list": result_scores_list,
|
|||
|
|
"result_classid_list": result_classid_list
|
|||
|
|
}
|
|||
|
|
# "box_detail": {
|
|||
|
|
# "result_box_list": result_box_list.tolist() if isinstance(result_box_list,
|
|||
|
|
# np.ndarray) else result_box_list,
|
|||
|
|
# "result_scores_list": result_scores_list.tolist() if isinstance(result_scores_list,
|
|||
|
|
# np.ndarray) else result_scores_list,
|
|||
|
|
# "result_classid_list": result_classid_list.tolist() if isinstance(result_classid_list,
|
|||
|
|
# np.ndarray) else result_classid_list
|
|||
|
|
# }
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
json_message = json.dumps(message, indent=4, ensure_ascii=False)
|
|||
|
|
mqtt_client.publish_message(json_message)
|
|||
|
|
time.sleep(0.1)
|
|||
|
|
if os.path.exists(minio_path):
|
|||
|
|
os.remove(minio_path)
|
|||
|
|
if os.path.exists(minio_path_before):
|
|||
|
|
os.remove(minio_path_before)
|
|||
|
|
except Exception as infer_err:
|
|||
|
|
print(f"推理错误: {infer_err}")
|
|||
|
|
finally:
|
|||
|
|
yolo11_wrapper.destroy()
|
|||
|
|
except Exception as infer_err:
|
|||
|
|
print(f"pic_detect 推理错误: {infer_err}")
|