396 lines
11 KiB
Python
396 lines
11 KiB
Python
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() |