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