115 lines
4.0 KiB
Python
115 lines
4.0 KiB
Python
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)
|
||
|