77 lines
2.4 KiB
Python
77 lines
2.4 KiB
Python
|
import os
|
|||
|
|
import shutil
|
|||
|
|
import cv2
|
|||
|
|
import collections
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
from AI_Python.Ai_tottle.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}")
|
|||
|
|
|
|||
|
|
# 检测图像(seg)
|
|||
|
|
|
|||
|
|
|
|||
|
|
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
|
|||
|
|
}
|
|||
|
|
|