AI_python_yoooger/Ai_tottle/map_find.py

import os
import shutil
import uuid
from miniohelp import upload_folder, load_config, parse_minio_url, download_file_url
from datetime import timedelta
from map.uav_seg import UAVSegPredictor
import errno

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 test_uav_seg(model_path, input_path, output_path):
    predictor = UAVSegPredictor(
        model_path= model_path,
        model_type='deeplabv3plus',
        num_classes=7
    )
    color_map = [
                [255, 0, 0],    # 类别0: 红色
                [0, 255, 0],    # 类别1: 绿色
                [0, 0, 255],    # 类别2: 蓝色
                [255, 255, 0],  # 类别3: 黄色
                [255, 0, 255],  # 类别4: 品红
                [0, 255, 255],  # 类别5: 青色
                [128, 0, 128]   # 类别6: 紫色
            ]
    predictor.predict_folder(input_path, output_path,color_map)
    
def map_process_images(urls, yaml_name, bucket_name, bucket_directory,uav_model_path):
    """
    批量处理图片：下载 -> 分割识别 -> 上传 MinIO -> 返回访问 URL 列表。
    
    参数:
    - urls: 图片 URL 列表（List[str]）
    - yaml_name: MinIO 配置文件路径
    - bucket_name: 存储桶名
    - bucket_directory: 存储桶目录（可共享一个目录）
    
    返回:
    - dict: 包含 code、msg 和上传后的 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 = []

        for url in urls:
            # 每个文件用唯一 ID 命名，避免冲突
            unique_name = str(uuid.uuid4()) + os.path.splitext(url)[-1]
            local_path = os.path.join(input_folder, unique_name)

            print(f"📥 下载图片: {url}")
            download_file_url(url, save_path=local_path)

        print("🧠 图像识别中...")
        test_uav_seg(uav_model_path, input_folder, output_path=output_folder)

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

        # 获取输出文件夹中的文件名并生成 MinIO URL
        result_files = list_files(output_folder)
        for result_path in result_files:
            filename = os.path.basename(result_path)
            object_path = os.path.join(bucket_directory, filename).replace("\\", "/")
            url = client.presigned_get_object(bucket_name, object_path, expires=timedelta(days=2))
            output_urls.append(parse_minio_url(url))

        # 清理临时文件夹
        shutil.rmtree(input_folder)
        shutil.rmtree(output_folder)

        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)

if __name__ == '__main__':
    urls = [
        "http://171.221.233.49:9000/300bdf2b-a150-406e-be63-d28bd29b409f/qiepian/data7/data3/15/25807/19269.png"]
    yaml_name = "config"
    bucket_name = "300bdf2b-a150-406e-be63-d28bd29b409f"
    bucket_directory = "2025/seg"
    uav_model_path = "D:/work/AI_Python/Ai_tottle/map/checkpoints/deeplabv3plus_best.pth"
    
    result = map_process_images(urls, yaml_name, bucket_name, bucket_directory,uav_model_path)
    print(result)