101 lines
3.7 KiB
Python
101 lines
3.7 KiB
Python
|
import os
|
|||
|
|
import shutil
|
|||
|
|
import cv2
|
|||
|
|
import collections
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
from miniohelp import downFile, upload_file, parse_minio_url # 确保你有这些工具函数
|
|||
|
|
from minio import Minio
|
|||
|
|
|
|||
|
|
def process_images(yolo_model, image_list, class_filter, input_folder, output_folder, minio_info):
|
|||
|
|
# 初始化 MinIO 客户端# 用配置字典初始化 Minio 客户端对象
|
|||
|
|
# 清洗 endpoint，去掉 http:// 或 https:// 前缀
|
|||
|
|
endpoint = minio_info["MinIOEndpoint"].replace("http://", "").replace("https://", "")
|
|||
|
|
|
|||
|
|
# 初始化 MinIO 客户端
|
|||
|
|
minio = Minio(
|
|||
|
|
endpoint=endpoint,
|
|||
|
|
access_key=minio_info["MinIOAccessKey"],
|
|||
|
|
secret_key=minio_info["MinIOSecretKey"],
|
|||
|
|
secure=False
|
|||
|
|
)
|
|||
|
|
os.makedirs(input_folder, exist_ok=True)
|
|||
|
|
os.makedirs(output_folder, exist_ok=True)
|
|||
|
|
|
|||
|
|
model = YOLO(yolo_model)
|
|||
|
|
class_ids_filter = [int(cls) for cls in class_filter.split(",")] if class_filter else None
|
|||
|
|
output_image_list = []
|
|||
|
|
|
|||
|
|
for item in image_list:
|
|||
|
|
img_id = item["id"]
|
|||
|
|
img_url = item["path"]
|
|||
|
|
|
|||
|
|
# 解析 MinIO 地址
|
|||
|
|
if img_url.startswith("http"):
|
|||
|
|
bucket_name, img_path = parse_minio_url(img_url)
|
|||
|
|
else:
|
|||
|
|
bucket_name, img_path = "default-bucket", img_url
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 下载原图到本地
|
|||
|
|
local_input_path = os.path.join(input_folder, os.path.basename(img_path))
|
|||
|
|
downFile(minio, img_path, bucket_name, local_input_path)
|
|||
|
|
|
|||
|
|
# 读取图像
|
|||
|
|
image = cv2.imread(local_input_path)
|
|||
|
|
if image is None:
|
|||
|
|
raise ValueError(f"无法读取图像: {local_input_path}")
|
|||
|
|
|
|||
|
|
# YOLO 检测
|
|||
|
|
results = model.predict(image,
|
|||
|
|
classes=class_ids_filter,
|
|||
|
|
conf=0.5,
|
|||
|
|
iou = 0.111,
|
|||
|
|
show_labels = False,)
|
|||
|
|
result = results[0]
|
|||
|
|
|
|||
|
|
# 统计类别数
|
|||
|
|
class_counts = collections.Counter(result.boxes.cls.cpu().numpy().astype(int)) if result.boxes is not None else {}
|
|||
|
|
filtered_class_counts = {k: v for k, v in class_counts.items() if k in class_ids_filter}
|
|||
|
|
|
|||
|
|
# 转换所有的 numpy.int64 为 Python 的 int 类型
|
|||
|
|
detected_classes = [int(cls) for cls in filtered_class_counts.keys()]
|
|||
|
|
detected_numbers = [int(num) for num in filtered_class_counts.values()]
|
|||
|
|
aim = bool(detected_classes)
|
|||
|
|
|
|||
|
|
# 保存标注图像
|
|||
|
|
annotated_image = result.plot(labels=False)
|
|||
|
|
filename_no_ext, ext = os.path.splitext(os.path.basename(img_path))
|
|||
|
|
output_filename = f"{filename_no_ext}_ai{ext}"
|
|||
|
|
local_output_path = os.path.join(output_folder, output_filename)
|
|||
|
|
cv2.imwrite(local_output_path, annotated_image)
|
|||
|
|
|
|||
|
|
# 上传标注图像到 MinIO
|
|||
|
|
minio_path = upload_file(minio, local_output_path, bucket_name, os.path.dirname(img_path))
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[错误] 处理失败 - {img_path}，错误: {str(e)}")
|
|||
|
|
detected_classes = []
|
|||
|
|
detected_numbers = []
|
|||
|
|
aim = False
|
|||
|
|
output_filename = ""
|
|||
|
|
minio_path = ""
|
|||
|
|
|
|||
|
|
output_image_list.append({
|
|||
|
|
"id": img_id,
|
|||
|
|
"minio_path":minio_path,
|
|||
|
|
"aim": aim,
|
|||
|
|
"class": detected_classes,
|
|||
|
|
"number": detected_numbers
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 清理临时目录
|
|||
|
|
shutil.rmtree(input_folder, ignore_errors=True)
|
|||
|
|
shutil.rmtree(output_folder, ignore_errors=True)
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"status": "success",
|
|||
|
|
"message": "Detection completed",
|
|||
|
|
"data": output_image_list
|
|||
|
|
}
|
|||
|
|
|