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