initial commit

app/image_pipeline.py

@@ -0,0 +1,121 @@
+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),
+    }