AI_python_yoooger/Ai_tottle/yolo_photo.py

import os
import shutil
import uuid
import errno
from ultralytics import YOLO
from miniohelp import upload_folder,load_config,download_file_url,parse_minio_url
from datetime import timedelta
def list_files(folder_path):
    """获取指定目录下的所有文件完整路径"""
    return [os.path.join(folder_path, f) for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f))]

def safe_rmtree(path):
    try:
        if os.path.exists(path):
            shutil.rmtree(path)
    except OSError as e:
        if e.errno != errno.ENOENT:  # 忽略“文件夹不存在”错误
            print(f"⚠️ 删除目录失败: {path} 错误: {e}")

def map_process_images(urls, yaml_name, bucket_name, bucket_directory, uav_model_path):
    """
    批量处理图片：下载 -> 分割识别 -> 上传 MinIO -> 返回访问 URL 列表。
    """
    input_folder = "input_folder"
    output_folder = "output_folder"
    os.makedirs(input_folder, exist_ok=True)
    os.makedirs(output_folder, exist_ok=True)

    try:
        print("🔐 加载 MinIO 配置...")
        client = load_config(yaml_name)
        output_urls = []

        # 下载图片到 input_folder
        for url in urls:
            unique_name = str(uuid.uuid4()) + os.path.splitext(url)[-1]
            local_path = os.path.join(input_folder, unique_name)
            print(f"📥 下载图片: {url} -> {local_path}")
            download_file_url(url, save_path=local_path)

        # YOLO 模型推理
        print("🧠 使用 YOLO 进行图像分割/识别...")
        model = YOLO(uav_model_path)
        model.predict(source=input_folder, save=True, project=output_folder, name="result", exist_ok=True)
        result_folder = os.path.join(output_folder, "result")  # 获取结果文件夹

        # 上传结果到 MinIO
        print("☁️ 上传结果到 MinIO...")
        upload_folder(client, result_folder, bucket_name, bucket_directory)

        # 生成访问 URL
        result_files = list_files(result_folder)
        for result_path in result_files:
            filename = os.path.basename(result_path)
            object_path = os.path.join(bucket_directory, filename).replace("\\", "/")
            presigned_url = client.presigned_get_object(bucket_name, object_path, expires=timedelta(days=2))
            output_urls.append(parse_minio_url(presigned_url))

        return {
            "code": 200,
            "msg": "success",
            "data": output_urls
        }

    except Exception as e:
        print(f"❌ 处理失败: {e}")
        return {
            "code": 500,
            "msg": f"error: {e}",
            "data": []
        }
    finally:
        safe_rmtree(input_folder)
        safe_rmtree(output_folder)

def map_process_images_with_progress(urls, yaml_name, bucket_name, bucket_directory, uav_model_path, progress_callback):
    """
    包装你的 map_process_images，支持进度回调。
    """

    input_folder = "input_folder"
    output_folder = "output_folder"
    os.makedirs(input_folder, exist_ok=True)
    os.makedirs(output_folder, exist_ok=True)

    try:
        client = load_config(yaml_name)
        output_urls = []
        total = len(urls)
        for i, url in enumerate(urls):
            unique_name = str(uuid.uuid4()) + os.path.splitext(url)[-1]
            local_path = os.path.join(input_folder, unique_name)
            download_file_url(url, save_path=local_path)
            progress_callback("downloading", int((i + 1) / total * 30))

        model = YOLO(uav_model_path)
        model.predict(source=input_folder, save=True, project=output_folder, name="result", exist_ok=True)
        result_folder = os.path.join(output_folder, "result")
        progress_callback("processing", 60)

        upload_folder(client, result_folder, bucket_name, bucket_directory)
        result_files = list_files(result_folder)
        for i, result_path in enumerate(result_files):
            filename = os.path.basename(result_path)
            object_path = os.path.join(bucket_directory, filename).replace("\\", "/")
            presigned_url = client.presigned_get_object(bucket_name, object_path)
            output_urls.append(parse_minio_url(presigned_url))
            progress_callback("uploading", 60 + int((i + 1) / len(result_files) * 40))

        return {"code": 200, "msg": "success", "data": output_urls}
    finally:
        safe_rmtree(input_folder)
        safe_rmtree(output_folder)