121 lines
3.4 KiB
Python
121 lines
3.4 KiB
Python
from pathlib import Path
|
|
from typing import Dict
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from app.config import settings
|
|
from app.utils import (
|
|
draw_results,
|
|
now_timestamp_minute,
|
|
safe_stem,
|
|
save_json,
|
|
write_image_xml,
|
|
)
|
|
|
|
|
|
def make_unique_result_dir(parent_dir: Path, timestamp: str, original_filename: str) -> tuple[Path, str, str]:
|
|
"""
|
|
创建单次图片识别结果目录。
|
|
|
|
目录名：
|
|
时间戳_原文件名stem
|
|
|
|
例如：
|
|
runs/images/202605181431_DJI_20250820151801_0006_Z/
|
|
|
|
返回：
|
|
result_dir, result_base_name, original_stem
|
|
"""
|
|
original_stem = safe_stem(original_filename)
|
|
base_name = f"{timestamp}_{original_stem}"
|
|
|
|
result_dir = parent_dir / base_name
|
|
|
|
if not result_dir.exists():
|
|
result_dir.mkdir(parents=True, exist_ok=True)
|
|
return result_dir, base_name, original_stem
|
|
|
|
idx = 1
|
|
while True:
|
|
candidate_base = f"{base_name}_{idx:02d}"
|
|
candidate_dir = parent_dir / candidate_base
|
|
if not candidate_dir.exists():
|
|
candidate_dir.mkdir(parents=True, exist_ok=True)
|
|
return candidate_dir, candidate_base, original_stem
|
|
idx += 1
|
|
|
|
|
|
def save_image_outputs(
|
|
img_bgr: np.ndarray,
|
|
original_filename: str,
|
|
predict_output: Dict,
|
|
) -> Dict:
|
|
"""
|
|
图片结果自动保存到：
|
|
runs/images/时间戳_原图片名/
|
|
|
|
目录内部生成：
|
|
原图片名.jpg
|
|
时间戳_原图片名.json
|
|
时间戳_原图片名.xml
|
|
|
|
示例：
|
|
runs/images/202605181431_DJI_20250820151801_0006_Z/
|
|
├── DJI_20250820151801_0006_Z.jpg
|
|
├── 202605181431_DJI_20250820151801_0006_Z.json
|
|
└── 202605181431_DJI_20250820151801_0006_Z.xml
|
|
"""
|
|
images_root = Path(settings.output_dir) / settings.image_output_subdir
|
|
images_root.mkdir(parents=True, exist_ok=True)
|
|
|
|
timestamp = now_timestamp_minute()
|
|
result_dir, result_base_name, original_stem = make_unique_result_dir(
|
|
parent_dir=images_root,
|
|
timestamp=timestamp,
|
|
original_filename=original_filename,
|
|
)
|
|
|
|
# 图片文件名不带时间戳，按你的要求保留原图片名 stem
|
|
image_path = result_dir / f"{original_stem}.jpg"
|
|
|
|
# JSON 和 XML 带时间戳
|
|
json_path = result_dir / f"{result_base_name}.json"
|
|
xml_path = result_dir / f"{result_base_name}.xml"
|
|
|
|
results = predict_output.get("results", [])
|
|
|
|
annotated = draw_results(img_bgr, results)
|
|
cv2.imwrite(str(image_path), annotated)
|
|
|
|
json_data = {
|
|
"success": True,
|
|
"type": "image",
|
|
"original_filename": original_filename,
|
|
"result_dir": str(result_dir),
|
|
"saved_image": str(image_path),
|
|
"saved_json": str(json_path),
|
|
"saved_xml": str(xml_path),
|
|
"image_width": predict_output.get("image_width"),
|
|
"image_height": predict_output.get("image_height"),
|
|
"num_objects": predict_output.get("num_objects"),
|
|
"results": results,
|
|
}
|
|
|
|
save_json(json_path, json_data)
|
|
|
|
write_image_xml(
|
|
xml_path=xml_path,
|
|
image_filename=image_path.name,
|
|
image_path=str(image_path),
|
|
width=int(predict_output.get("image_width", img_bgr.shape[1])),
|
|
height=int(predict_output.get("image_height", img_bgr.shape[0])),
|
|
results=results,
|
|
)
|
|
|
|
return {
|
|
"result_dir": str(result_dir),
|
|
"image": str(image_path),
|
|
"json": str(json_path),
|
|
"xml": str(xml_path),
|
|
} |