from __future__ import annotations import argparse import json import time from pathlib import Path from typing import Dict, List import cv2 from app.config import settings from app.pipeline import BirdCascadePipeline from app.utils import ( decode_image_from_bytes, draw_results, safe_stem, save_json, write_image_xml, write_video_xml, ) from app.video_pipeline import VideoResultSmoother IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"} VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv", ".flv", ".wmv", ".m4v"} def now_timestamp_minute() -> str: return time.strftime("%Y%m%d%H%M") def is_image(path: Path) -> bool: return path.suffix.lower() in IMAGE_EXTS def is_video(path: Path) -> bool: return path.suffix.lower() in VIDEO_EXTS def collect_inputs(input_path: Path, recursive: bool = False) -> List[Path]: if input_path.is_file(): if is_image(input_path) or is_video(input_path): return [input_path] raise ValueError(f"不支持的文件类型: {input_path}") if not input_path.is_dir(): raise FileNotFoundError(f"输入路径不存在: {input_path}") iterator = input_path.rglob("*") if recursive else input_path.glob("*") files = [] for p in iterator: if p.is_file() and (is_image(p) or is_video(p)): files.append(p) files.sort() return files def make_output_dirs(input_path: Path, output_base: Path) -> Dict[str, Path]: timestamp = now_timestamp_minute() input_name = safe_stem(input_path.name) output_root = output_base / f"{timestamp}_{input_name}" if output_root.exists(): idx = 1 while True: candidate = output_base / f"{timestamp}_{input_name}_{idx:02d}" if not candidate.exists(): output_root = candidate break idx += 1 result_dir = output_root / "result" xml_dir = output_root / "xml" json_dir = output_root / "json" result_dir.mkdir(parents=True, exist_ok=True) xml_dir.mkdir(parents=True, exist_ok=True) json_dir.mkdir(parents=True, exist_ok=True) return { "root": output_root, "result": result_dir, "xml": xml_dir, "json": json_dir, } def unique_path(path: Path) -> Path: """ 如果不同子目录里存在同名文件，避免覆盖。 例如 result/1.jpg 已存在，则自动变成 result/1_01.jpg。 """ if not path.exists(): return path stem = path.stem suffix = path.suffix parent = path.parent idx = 1 while True: candidate = parent / f"{stem}_{idx:02d}{suffix}" if not candidate.exists(): return candidate idx += 1 def process_image( image_path: Path, pipeline: BirdCascadePipeline, dirs: Dict[str, Path], ) -> Dict: img_bgr = decode_image_from_bytes(image_path.read_bytes()) output = pipeline.predict(img_bgr, return_image=False) results = output["results"] result_img_path = unique_path(dirs["result"] / f"{image_path.stem}.jpg") json_path = unique_path(dirs["json"] / f"{image_path.stem}.json") xml_path = unique_path(dirs["xml"] / f"{image_path.stem}.xml") annotated = draw_results(img_bgr, results) cv2.imwrite(str(result_img_path), annotated) json_data = { "success": True, "type": "image", "source_path": str(image_path), "result_path": str(result_img_path), "xml_path": str(xml_path), "json_path": str(json_path), "image_width": output["image_width"], "image_height": output["image_height"], "num_objects": output["num_objects"], "results": results, } save_json(json_path, json_data) write_image_xml( xml_path=xml_path, image_filename=result_img_path.name, image_path=str(result_img_path), width=int(output["image_width"]), height=int(output["image_height"]), results=results, ) return json_data def process_video( video_path: Path, pipeline: BirdCascadePipeline, dirs: Dict[str, Path], sample_interval: int, ) -> Dict: result_video_path = unique_path(dirs["result"] / f"{video_path.stem}.mp4") json_path = unique_path(dirs["json"] / f"{video_path.stem}.json") xml_path = unique_path(dirs["xml"] / f"{video_path.stem}.xml") cap = cv2.VideoCapture(str(video_path)) if not cap.isOpened(): raise ValueError(f"视频打开失败: {video_path}") fps = cap.get(cv2.CAP_PROP_FPS) if fps is None or fps <= 0: fps = 25.0 width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0) fourcc = cv2.VideoWriter_fourcc(*"mp4v") writer = cv2.VideoWriter(str(result_video_path), fourcc, fps, (width, height)) smoother = VideoResultSmoother( hold_frames=settings.smooth_hold_frames, alpha=settings.box_smooth_alpha, ) frame_results = [] processed_frames = 0 detected_frames = 0 frame_index = 0 start_time = time.time() try: while True: ok, frame = cap.read() if not ok: break should_process = frame_index % sample_interval == 0 if should_process: output = pipeline.predict(frame, return_image=False) raw_results = output["results"] processed_frames += 1 if settings.smooth_video_boxes: results = smoother.update(raw_results, frame_index) else: results = raw_results if len(results) > 0: detected_frames += 1 frame_results.append( { "frame_index": frame_index, "time_sec": round(frame_index / fps, 3), "num_objects": len(results), "results": results, } ) else: if settings.smooth_video_boxes: results = smoother.get_for_skipped_frame(frame_index) else: results = [] annotated_frame = draw_results(frame, results) writer.write(annotated_frame) frame_index += 1 finally: cap.release() writer.release() elapsed = round(time.time() - start_time, 3) summary = { "success": True, "type": "video", "source_path": str(video_path), "result_path": str(result_video_path), "xml_path": str(xml_path), "json_path": str(json_path), "video_info": { "width": width, "height": height, "fps": fps, "total_frames": total_frames, }, "sample_interval": sample_interval, "smooth_video_boxes": settings.smooth_video_boxes, "smooth_hold_frames": settings.smooth_hold_frames, "box_smooth_alpha": settings.box_smooth_alpha, "processed_frames": processed_frames, "detected_frames": detected_frames, "elapsed_sec": elapsed, } json_data = { "summary": summary, "frames": frame_results, } save_json(json_path, json_data) write_video_xml( xml_path=xml_path, video_filename=result_video_path.name, video_path=str(result_video_path), width=width, height=height, fps=fps, total_frames=total_frames, sample_interval=sample_interval, frame_results=frame_results, ) return summary def run_batch( input_path: str, output_base: str, recursive: bool, sample_interval: int, ) -> Dict: input_path_obj = Path(input_path).resolve() output_base_obj = Path(output_base).resolve() output_base_obj.mkdir(parents=True, exist_ok=True) dirs = make_output_dirs(input_path_obj, output_base_obj) files = collect_inputs(input_path_obj, recursive=recursive) if len(files) == 0: raise ValueError(f"没有找到图片或视频文件: {input_path_obj}") pipeline = BirdCascadePipeline() success_results = [] failed_results = [] for idx, file_path in enumerate(files, start=1): print(f"[{idx}/{len(files)}] processing: {file_path}") try: if is_image(file_path): result = process_image( image_path=file_path, pipeline=pipeline, dirs=dirs, ) elif is_video(file_path): result = process_video( video_path=file_path, pipeline=pipeline, dirs=dirs, sample_interval=sample_interval, ) else: continue success_results.append(result) except Exception as e: failed = { "source_path": str(file_path), "error": str(e), } failed_results.append(failed) print(f"[ERROR] {file_path}: {e}") index_json = { "success": len(failed_results) == 0, "input_path": str(input_path_obj), "output_root": str(dirs["root"]), "result_dir": str(dirs["result"]), "xml_dir": str(dirs["xml"]), "json_dir": str(dirs["json"]), "total_files": len(files), "success_files": len(success_results), "failed_files": len(failed_results), "failed": failed_results, } index_path = dirs["root"] / "index.json" index_path.write_text( json.dumps(index_json, ensure_ascii=False, indent=2), encoding="utf-8", ) print("\nDone.") print(f"Output root: {dirs['root']}") print(f"Result dir : {dirs['result']}") print(f"XML dir : {dirs['xml']}") print(f"JSON dir : {dirs['json']}") print(f"Index JSON : {index_path}") return index_json def main(): parser = argparse.ArgumentParser(description="Bird cascade batch inference") parser.add_argument( "--input", required=True, help="输入路径，可以是单张图片、单个视频，或包含图片/视频的文件夹", ) parser.add_argument( "--output", default="/data2/zhangw_data/Bird_detect_cascade_service/runs", help="输出根目录，默认直接输出到项目 runs 下", ) parser.add_argument( "--recursive", action="store_true", help="输入为文件夹时，是否递归处理子文件夹", ) parser.add_argument( "--sample-interval", type=int, default=settings.video_sample_interval, help="视频抽帧间隔", ) args = parser.parse_args() run_batch( input_path=args.input, output_base=args.output, recursive=args.recursive, sample_interval=args.sample_interval, ) if __name__ == "__main__": main()