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
|
||
}
|
||
|