initial commit
This commit is contained in:
Ankle committed
commit
0bafbee31f
54 files changed
+6317
No files matched your search
@@ -0,0 +1,396 @@
|
||||
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()
|
||||
Reference in new issue
Block a user