925 lines
30 KiB
Python
925 lines
30 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Dict, List, Sequence
|
|
|
|
import cv2
|
|
|
|
from app.config import settings
|
|
from app.pipeline import BirdCascadePipeline
|
|
from app.utils import (
|
|
decode_image_from_bytes,
|
|
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 parse_bool(value: str | bool) -> bool:
|
|
if isinstance(value, bool):
|
|
return value
|
|
|
|
normalized = value.strip().lower()
|
|
if normalized in {"1", "true", "yes", "y", "on"}:
|
|
return True
|
|
if normalized in {"0", "false", "no", "n", "off"}:
|
|
return False
|
|
|
|
raise argparse.ArgumentTypeError(
|
|
f"布尔值必须是 true/false、yes/no、1/0 之一，当前为: {value}"
|
|
)
|
|
|
|
|
|
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:
|
|
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"]
|
|
density_json = output.get("density")
|
|
density_result = output.get("density_result")
|
|
|
|
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 = pipeline.draw_full_result(
|
|
img_bgr=img_bgr,
|
|
results=results,
|
|
density_result=density_result,
|
|
)
|
|
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"],
|
|
"density": density_json,
|
|
"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()
|
|
|
|
last_density_result = None
|
|
last_density_json = None
|
|
|
|
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"]
|
|
density_result = output.get("density_result")
|
|
density_json = output.get("density")
|
|
|
|
last_density_result = density_result
|
|
last_density_json = density_json
|
|
|
|
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),
|
|
"density": density_json,
|
|
"results": results,
|
|
}
|
|
)
|
|
|
|
else:
|
|
if settings.smooth_video_boxes:
|
|
results = smoother.get_for_skipped_frame(frame_index)
|
|
else:
|
|
results = []
|
|
|
|
density_result = last_density_result
|
|
density_json = last_density_json
|
|
|
|
annotated_frame = pipeline.draw_full_result(
|
|
img_bgr=frame,
|
|
results=results,
|
|
density_result=density_result,
|
|
)
|
|
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 build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description="Bird cascade batch inference")
|
|
parser.add_argument(
|
|
"--input",
|
|
required=True,
|
|
help="输入路径，可以是单张图片、单个视频，或包含图片/视频的文件夹",
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
default=settings.output_dir,
|
|
help="输出根目录，默认输出到项目 runs 下",
|
|
)
|
|
|
|
recursive_group = parser.add_mutually_exclusive_group()
|
|
recursive_group.add_argument(
|
|
"--recursive",
|
|
dest="recursive",
|
|
action="store_true",
|
|
help="输入为文件夹时递归处理子文件夹（默认）",
|
|
)
|
|
recursive_group.add_argument(
|
|
"--no-recursive",
|
|
dest="recursive",
|
|
action="store_false",
|
|
help="输入为文件夹时只处理第一层文件",
|
|
)
|
|
parser.set_defaults(recursive=True)
|
|
|
|
parser.add_argument(
|
|
"--sample-interval",
|
|
type=int,
|
|
default=settings.video_sample_interval,
|
|
help="视频抽帧间隔",
|
|
)
|
|
|
|
detection_group = parser.add_mutually_exclusive_group()
|
|
detection_group.add_argument(
|
|
"--enable-detection",
|
|
dest="enable_detection",
|
|
action="store_true",
|
|
help="启用 YOLO 目标检测（默认）",
|
|
)
|
|
detection_group.add_argument(
|
|
"--disable-detection",
|
|
dest="enable_detection",
|
|
action="store_false",
|
|
help="关闭 YOLO 目标检测",
|
|
)
|
|
parser.set_defaults(enable_detection=settings.enable_detection)
|
|
|
|
classification_group = parser.add_mutually_exclusive_group()
|
|
classification_group.add_argument(
|
|
"--enable-classification",
|
|
dest="enable_classification",
|
|
action="store_true",
|
|
help="启用 ConvNeXt 细粒度识别（默认）",
|
|
)
|
|
classification_group.add_argument(
|
|
"--disable-classification",
|
|
dest="enable_classification",
|
|
action="store_false",
|
|
help="关闭 ConvNeXt 细粒度识别，仅保留 YOLO 检测框",
|
|
)
|
|
parser.set_defaults(enable_classification=settings.enable_classification)
|
|
|
|
density_group = parser.add_mutually_exclusive_group()
|
|
density_group.add_argument(
|
|
"--enable-density",
|
|
dest="enable_density",
|
|
action="store_true",
|
|
help="启用贝叶斯密度估计（默认）",
|
|
)
|
|
density_group.add_argument(
|
|
"--disable-density",
|
|
dest="enable_density",
|
|
action="store_false",
|
|
help="关闭贝叶斯密度估计",
|
|
)
|
|
parser.set_defaults(enable_density=settings.enable_density)
|
|
|
|
parser.add_argument(
|
|
"--det-conf-thres",
|
|
type=float,
|
|
default=settings.det_conf_thres,
|
|
help="YOLO 第一层检测置信度阈值",
|
|
)
|
|
parser.add_argument(
|
|
"--det-iou-thres",
|
|
type=float,
|
|
default=settings.det_iou_thres,
|
|
help="NMS IoU 阈值",
|
|
)
|
|
parser.add_argument(
|
|
"--det-second-conf-thres",
|
|
type=float,
|
|
default=settings.det_second_conf_thres,
|
|
help="进入 ConvNeXt 前的第二层检测置信度阈值",
|
|
)
|
|
parser.add_argument(
|
|
"--min-box-area-ratio",
|
|
type=float,
|
|
default=settings.min_box_area_ratio,
|
|
help="检测框最小面积占比",
|
|
)
|
|
parser.add_argument(
|
|
"--max-box-area-ratio",
|
|
type=float,
|
|
default=settings.max_box_area_ratio,
|
|
help="检测框最大面积占比",
|
|
)
|
|
parser.add_argument(
|
|
"--max-box-aspect-ratio",
|
|
type=float,
|
|
default=settings.max_box_aspect_ratio,
|
|
help="检测框最大长宽比",
|
|
)
|
|
parser.add_argument(
|
|
"--detection-roi-y-min-ratio",
|
|
type=float,
|
|
default=settings.detection_roi_y_min_ratio,
|
|
help="只保留中心点位于该纵向比例以下区域之外的框，默认 0",
|
|
)
|
|
parser.add_argument(
|
|
"--detection-roi-y-max-ratio",
|
|
type=float,
|
|
default=settings.detection_roi_y_max_ratio,
|
|
help="只保留中心点位于该纵向比例以上区域之前的框，默认 1",
|
|
)
|
|
parser.add_argument(
|
|
"--cls-show-conf-thres",
|
|
type=float,
|
|
default=settings.cls_show_conf_thres,
|
|
help="普通目标显示物种名的分类置信度阈值",
|
|
)
|
|
parser.add_argument(
|
|
"--large-box-area-ratio",
|
|
type=float,
|
|
default=settings.large_box_area_ratio,
|
|
help="近景大目标的框面积占比阈值",
|
|
)
|
|
parser.add_argument(
|
|
"--large-box-cls-show-conf-thres",
|
|
type=float,
|
|
default=settings.large_box_cls_show_conf_thres,
|
|
help="大目标显示物种名的分类置信度阈值",
|
|
)
|
|
parser.add_argument(
|
|
"--show-large-box-name-always",
|
|
type=parse_bool,
|
|
default=settings.show_large_box_name_always,
|
|
metavar="BOOL",
|
|
help="大目标是否始终显示非空物种名",
|
|
)
|
|
parser.add_argument(
|
|
"--dense-scene-count-thres",
|
|
type=int,
|
|
default=settings.dense_scene_count_thres,
|
|
help="判定密集场景的密度估计数量阈值",
|
|
)
|
|
parser.add_argument(
|
|
"--hide-cls-name-in-dense-scene",
|
|
type=parse_bool,
|
|
default=settings.hide_cls_name_in_dense_scene,
|
|
metavar="BOOL",
|
|
help="密集场景下是否隐藏物种名",
|
|
)
|
|
parser.add_argument(
|
|
"--crop-padding-ratio",
|
|
type=float,
|
|
default=settings.crop_padding_ratio,
|
|
help="送入 ConvNeXt 前检测框向外扩的比例",
|
|
)
|
|
parser.add_argument(
|
|
"--max-detections",
|
|
type=int,
|
|
default=settings.max_detections,
|
|
help="单张图或单帧最多保留的检测框数量",
|
|
)
|
|
parser.add_argument(
|
|
"--uniform-species-color",
|
|
type=parse_bool,
|
|
default=settings.uniform_species_color,
|
|
metavar="BOOL",
|
|
help="所有物种的检测框和标签是否统一使用同一种颜色",
|
|
)
|
|
parser.add_argument(
|
|
"--force-species-name",
|
|
dest="forced_species_name",
|
|
default=settings.forced_species_name,
|
|
help="人工确认整段视频为单一物种时，强制覆盖所有细粒度标签；留空不覆盖",
|
|
)
|
|
parser.add_argument(
|
|
"--force-species-class-id",
|
|
dest="forced_species_class_id",
|
|
type=int,
|
|
default=settings.forced_species_class_id,
|
|
help="强制物种对应的类别 ID；默认 -1 时按标签文件自动查找",
|
|
)
|
|
parser.add_argument(
|
|
"--force-species-min-box-area-ratio",
|
|
dest="forced_species_min_box_area_ratio",
|
|
type=float,
|
|
default=settings.forced_species_min_box_area_ratio,
|
|
help="只覆盖面积占比达到该值的目标；默认 0 表示覆盖全部检测目标",
|
|
)
|
|
parser.add_argument(
|
|
"--species-name-remap",
|
|
default=settings.species_name_remap,
|
|
help="按人工复核纠正指定标签，格式：东方白鹳=白鹤,白鹭=白鹤",
|
|
)
|
|
parser.add_argument(
|
|
"--species-name-remap-min-confidence",
|
|
type=float,
|
|
default=settings.species_name_remap_min_confidence,
|
|
help="只有原始分类置信度达到该值时才应用物种标签映射",
|
|
)
|
|
parser.add_argument(
|
|
"--allowed-species-names",
|
|
default=settings.allowed_species_names,
|
|
help="只允许这些物种显示标签，英文逗号分隔；留空表示不限制",
|
|
)
|
|
parser.add_argument(
|
|
"--binary-dark-species-name",
|
|
default=settings.binary_dark_species_name,
|
|
help="双物种明暗规则中的深色物种名称；留空表示关闭",
|
|
)
|
|
parser.add_argument(
|
|
"--binary-light-species-name",
|
|
default=settings.binary_light_species_name,
|
|
help="双物种明暗规则中的浅色物种名称；留空表示关闭",
|
|
)
|
|
parser.add_argument(
|
|
"--binary-dark-pixel-threshold",
|
|
type=int,
|
|
default=settings.binary_dark_pixel_threshold,
|
|
help="灰度低于该值的像素视为暗色像素，范围 0~255",
|
|
)
|
|
parser.add_argument(
|
|
"--binary-dark-ratio-threshold",
|
|
type=float,
|
|
default=settings.binary_dark_ratio_threshold,
|
|
help="暗色像素比例达到该值时使用深色物种标签",
|
|
)
|
|
parser.add_argument(
|
|
"--dark-override-species-name",
|
|
default=settings.dark_override_species_name,
|
|
help="暗色目标人工覆盖标签；留空表示关闭",
|
|
)
|
|
parser.add_argument(
|
|
"--dark-override-species-class-id",
|
|
type=int,
|
|
default=settings.dark_override_species_class_id,
|
|
help="暗色目标人工覆盖标签的 class id；默认 -1 时按类别文件查找",
|
|
)
|
|
parser.add_argument(
|
|
"--dark-override-pixel-threshold",
|
|
type=int,
|
|
default=settings.dark_override_pixel_threshold,
|
|
help="灰度低于该值的像素视为暗色像素，范围 0~255",
|
|
)
|
|
parser.add_argument(
|
|
"--dark-override-ratio-threshold",
|
|
type=float,
|
|
default=settings.dark_override_ratio_threshold,
|
|
help="暗色像素比例达到该值时覆盖为指定标签",
|
|
)
|
|
parser.add_argument(
|
|
"--dark-override-min-box-area-ratio",
|
|
type=float,
|
|
default=settings.dark_override_min_box_area_ratio,
|
|
help="只覆盖面积占比达到该值的暗色目标",
|
|
)
|
|
parser.add_argument(
|
|
"--dark-override-max-box-area-ratio",
|
|
type=float,
|
|
default=settings.dark_override_max_box_area_ratio,
|
|
help="只覆盖面积占比不超过该值的暗色目标",
|
|
)
|
|
parser.add_argument(
|
|
"--dark-override-roi-y-min-ratio",
|
|
type=float,
|
|
default=settings.dark_override_roi_y_min_ratio,
|
|
help="暗色覆盖只作用于中心点纵向比例不小于该值的目标",
|
|
)
|
|
parser.add_argument(
|
|
"--dark-override-roi-y-max-ratio",
|
|
type=float,
|
|
default=settings.dark_override_roi_y_max_ratio,
|
|
help="暗色覆盖只作用于中心点纵向比例不大于该值的目标",
|
|
)
|
|
parser.add_argument(
|
|
"--region-override-species-name",
|
|
default=settings.region_override_species_name,
|
|
help="区域人工覆盖标签；留空表示关闭",
|
|
)
|
|
parser.add_argument(
|
|
"--region-override-species-class-id",
|
|
type=int,
|
|
default=settings.region_override_species_class_id,
|
|
help="区域人工覆盖标签的 class id；默认 -1 时按类别文件查找",
|
|
)
|
|
parser.add_argument(
|
|
"--region-override-source-species-names",
|
|
default=settings.region_override_source_species_names,
|
|
help="只覆盖这些原标签，英文逗号分隔；留空表示覆盖区域内全部非空标签",
|
|
)
|
|
parser.add_argument(
|
|
"--region-override-roi-y-min-ratio",
|
|
type=float,
|
|
default=settings.region_override_roi_y_min_ratio,
|
|
help="区域覆盖只作用于中心点纵向比例不小于该值的目标",
|
|
)
|
|
parser.add_argument(
|
|
"--region-override-roi-y-max-ratio",
|
|
type=float,
|
|
default=settings.region_override_roi_y_max_ratio,
|
|
help="区域覆盖只作用于中心点纵向比例不大于该值的目标",
|
|
)
|
|
parser.add_argument(
|
|
"--enable-crane-appearance-correction",
|
|
action="store_true",
|
|
default=settings.enable_crane_appearance_correction,
|
|
help="启用白鹤/白头鹤/白枕鹤外观二次纠正",
|
|
)
|
|
parser.add_argument(
|
|
"--crane-white-ratio-threshold",
|
|
type=float,
|
|
default=settings.crane_white_ratio_threshold,
|
|
help="鹤类纠正中判定白鹤的亮白像素比例阈值",
|
|
)
|
|
parser.add_argument(
|
|
"--crane-dark-ratio-threshold",
|
|
type=float,
|
|
default=settings.crane_dark_ratio_threshold,
|
|
help="鹤类纠正中判定白枕鹤/白头鹤的暗色像素比例阈值",
|
|
)
|
|
parser.add_argument(
|
|
"--crane-tail-white-ratio-threshold",
|
|
type=float,
|
|
default=settings.crane_tail_white_ratio_threshold,
|
|
help="鹤类纠正中判定白枕鹤白尾/白斑的局部亮白比例阈值",
|
|
)
|
|
parser.add_argument(
|
|
"--crane-uncertain-hide-label",
|
|
type=parse_bool,
|
|
default=settings.crane_uncertain_hide_label,
|
|
metavar="BOOL",
|
|
help="鹤类纠正无法确认时是否隐藏标签",
|
|
)
|
|
|
|
return parser
|
|
|
|
|
|
def validate_args(
|
|
parser: argparse.ArgumentParser,
|
|
args: argparse.Namespace,
|
|
) -> argparse.Namespace:
|
|
if args.sample_interval < 1:
|
|
parser.error("--sample-interval 必须大于等于 1")
|
|
|
|
probability_fields = {
|
|
"--det-conf-thres": args.det_conf_thres,
|
|
"--det-iou-thres": args.det_iou_thres,
|
|
"--det-second-conf-thres": args.det_second_conf_thres,
|
|
"--cls-show-conf-thres": args.cls_show_conf_thres,
|
|
"--large-box-cls-show-conf-thres": args.large_box_cls_show_conf_thres,
|
|
"--species-name-remap-min-confidence": (
|
|
args.species_name_remap_min_confidence
|
|
),
|
|
"--detection-roi-y-min-ratio": args.detection_roi_y_min_ratio,
|
|
"--detection-roi-y-max-ratio": args.detection_roi_y_max_ratio,
|
|
"--binary-dark-ratio-threshold": args.binary_dark_ratio_threshold,
|
|
"--dark-override-ratio-threshold": args.dark_override_ratio_threshold,
|
|
"--dark-override-roi-y-min-ratio": args.dark_override_roi_y_min_ratio,
|
|
"--dark-override-roi-y-max-ratio": args.dark_override_roi_y_max_ratio,
|
|
"--region-override-roi-y-min-ratio": (
|
|
args.region_override_roi_y_min_ratio
|
|
),
|
|
"--region-override-roi-y-max-ratio": (
|
|
args.region_override_roi_y_max_ratio
|
|
),
|
|
"--crane-white-ratio-threshold": args.crane_white_ratio_threshold,
|
|
"--crane-dark-ratio-threshold": args.crane_dark_ratio_threshold,
|
|
"--crane-tail-white-ratio-threshold": (
|
|
args.crane_tail_white_ratio_threshold
|
|
),
|
|
}
|
|
for option, value in probability_fields.items():
|
|
if not 0.0 <= value <= 1.0:
|
|
parser.error(f"{option} 必须在 0 到 1 之间")
|
|
|
|
ratio_fields = {
|
|
"--min-box-area-ratio": args.min_box_area_ratio,
|
|
"--max-box-area-ratio": args.max_box_area_ratio,
|
|
"--large-box-area-ratio": args.large_box_area_ratio,
|
|
"--crop-padding-ratio": args.crop_padding_ratio,
|
|
"--force-species-min-box-area-ratio": (
|
|
args.forced_species_min_box_area_ratio
|
|
),
|
|
"--dark-override-min-box-area-ratio": (
|
|
args.dark_override_min_box_area_ratio
|
|
),
|
|
"--dark-override-max-box-area-ratio": (
|
|
args.dark_override_max_box_area_ratio
|
|
),
|
|
}
|
|
for option, value in ratio_fields.items():
|
|
if value < 0.0:
|
|
parser.error(f"{option} 不能小于 0")
|
|
|
|
if args.min_box_area_ratio > args.max_box_area_ratio:
|
|
parser.error("--min-box-area-ratio 不能大于 --max-box-area-ratio")
|
|
if args.detection_roi_y_min_ratio > args.detection_roi_y_max_ratio:
|
|
parser.error(
|
|
"--detection-roi-y-min-ratio 不能大于 "
|
|
"--detection-roi-y-max-ratio"
|
|
)
|
|
if args.max_box_aspect_ratio <= 0.0:
|
|
parser.error("--max-box-aspect-ratio 必须大于 0")
|
|
if args.dense_scene_count_thres < 0:
|
|
parser.error("--dense-scene-count-thres 不能小于 0")
|
|
if args.max_detections < 1:
|
|
parser.error("--max-detections 必须大于等于 1")
|
|
if not 0 <= args.binary_dark_pixel_threshold <= 255:
|
|
parser.error("--binary-dark-pixel-threshold 必须在 0 到 255 之间")
|
|
if not 0 <= args.dark_override_pixel_threshold <= 255:
|
|
parser.error("--dark-override-pixel-threshold 必须在 0 到 255 之间")
|
|
if args.dark_override_min_box_area_ratio > args.dark_override_max_box_area_ratio:
|
|
parser.error(
|
|
"--dark-override-min-box-area-ratio 不能大于 "
|
|
"--dark-override-max-box-area-ratio"
|
|
)
|
|
if args.dark_override_roi_y_min_ratio > args.dark_override_roi_y_max_ratio:
|
|
parser.error(
|
|
"--dark-override-roi-y-min-ratio 不能大于 "
|
|
"--dark-override-roi-y-max-ratio"
|
|
)
|
|
if args.region_override_roi_y_min_ratio > args.region_override_roi_y_max_ratio:
|
|
parser.error(
|
|
"--region-override-roi-y-min-ratio 不能大于 "
|
|
"--region-override-roi-y-max-ratio"
|
|
)
|
|
if bool(args.binary_dark_species_name) != bool(
|
|
args.binary_light_species_name
|
|
):
|
|
parser.error(
|
|
"--binary-dark-species-name 和 "
|
|
"--binary-light-species-name 必须同时设置"
|
|
)
|
|
|
|
return args
|
|
|
|
|
|
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
|
parser = build_parser()
|
|
return validate_args(parser, parser.parse_args(argv))
|
|
|
|
|
|
def apply_cli_settings(args: argparse.Namespace) -> None:
|
|
setting_names = (
|
|
"enable_detection",
|
|
"enable_classification",
|
|
"enable_density",
|
|
"det_conf_thres",
|
|
"det_iou_thres",
|
|
"det_second_conf_thres",
|
|
"min_box_area_ratio",
|
|
"max_box_area_ratio",
|
|
"max_box_aspect_ratio",
|
|
"detection_roi_y_min_ratio",
|
|
"detection_roi_y_max_ratio",
|
|
"cls_show_conf_thres",
|
|
"large_box_area_ratio",
|
|
"large_box_cls_show_conf_thres",
|
|
"show_large_box_name_always",
|
|
"dense_scene_count_thres",
|
|
"hide_cls_name_in_dense_scene",
|
|
"crop_padding_ratio",
|
|
"max_detections",
|
|
"uniform_species_color",
|
|
"forced_species_name",
|
|
"forced_species_class_id",
|
|
"forced_species_min_box_area_ratio",
|
|
"species_name_remap",
|
|
"species_name_remap_min_confidence",
|
|
"allowed_species_names",
|
|
"binary_dark_species_name",
|
|
"binary_light_species_name",
|
|
"binary_dark_pixel_threshold",
|
|
"binary_dark_ratio_threshold",
|
|
"dark_override_species_name",
|
|
"dark_override_species_class_id",
|
|
"dark_override_pixel_threshold",
|
|
"dark_override_ratio_threshold",
|
|
"dark_override_min_box_area_ratio",
|
|
"dark_override_max_box_area_ratio",
|
|
"dark_override_roi_y_min_ratio",
|
|
"dark_override_roi_y_max_ratio",
|
|
"region_override_species_name",
|
|
"region_override_species_class_id",
|
|
"region_override_source_species_names",
|
|
"region_override_roi_y_min_ratio",
|
|
"region_override_roi_y_max_ratio",
|
|
"enable_crane_appearance_correction",
|
|
"crane_white_ratio_threshold",
|
|
"crane_dark_ratio_threshold",
|
|
"crane_tail_white_ratio_threshold",
|
|
"crane_uncertain_hide_label",
|
|
)
|
|
|
|
for name in setting_names:
|
|
object.__setattr__(settings, name, getattr(args, name))
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None):
|
|
args = parse_args(argv)
|
|
apply_cli_settings(args)
|
|
|
|
return run_batch(
|
|
input_path=args.input,
|
|
output_base=args.output,
|
|
recursive=args.recursive,
|
|
sample_interval=args.sample_interval,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|