jetson
07289a419a
- detect_interval=2 on both pipelines; stabilizer 10-frame hold bridges skipped frames (GPU inference cut ~half during active batch) - sup_det_interval 5 -> 10 (still <= stabilizer hold / synthetic max_age) - preview_every_n default 2 -> 5 - AnnotatedVideoWriter: codec auto-chain gstreamer_nvenc -> avc1 -> mp4v, max_height downscale (even dims, INTER_AREA); backend logged - output_max_height plumbed Job -> web checkbox (720p) -> CLI --output-height - README: 720p option, CLI flag - 153 tests pass (+13)
136 lines
5.2 KiB
Python
136 lines
5.2 KiB
Python
"""CLI entry point for feedmill_recounter (console script `recounter`)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
|
|
from src.model_registry import ModelConfig, scan_models
|
|
from src.pipeline import run_pipeline
|
|
|
|
|
|
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Feedmill Recounter — AI video analysis for object counting"
|
|
)
|
|
parser.add_argument("--video", type=str, default=None,
|
|
help="Input video file")
|
|
parser.add_argument("--models-dir", type=str, default="./models",
|
|
help="Directory containing model weight files")
|
|
parser.add_argument("--model", type=str, action="append", default=None,
|
|
help="Model filename to run (repeatable)")
|
|
parser.add_argument("--all-models", action="store_true",
|
|
help="Run all discovered models")
|
|
parser.add_argument("--list-models", action="store_true",
|
|
help="List discovered models and exit")
|
|
parser.add_argument("--filter", type=str, action="append", default=None,
|
|
help="Class name to keep (repeatable)")
|
|
parser.add_argument("--sack-conf", type=float, default=0.4,
|
|
help="Sack detection confidence threshold")
|
|
parser.add_argument("--truck-conf", type=float, default=0.5,
|
|
help="Truck detection confidence threshold")
|
|
parser.add_argument("--output", type=str, default=None,
|
|
help="Output path (honored only for single-model runs)")
|
|
parser.add_argument("--output-dir", type=str, default="./output",
|
|
help="Output directory for annotated videos")
|
|
parser.add_argument("--output-height", type=int, default=None,
|
|
help="Downscale annotated output to this height (e.g. 720)")
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> None:
|
|
args = parse_args(argv)
|
|
|
|
discovered = scan_models(args.models_dir)
|
|
|
|
if args.list_models:
|
|
if not discovered:
|
|
print(f"No models found in {args.models_dir}")
|
|
sys.exit(1)
|
|
print(f"{'Filename':<55} Classes")
|
|
for cfg in discovered:
|
|
classes = ", ".join(cfg.known_classes) if cfg.known_classes else "(unknown)"
|
|
print(f"{cfg.filename:<55} {classes}")
|
|
sys.exit(0)
|
|
|
|
if not args.video:
|
|
print("Error: --video is required (or use --list-models)", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
if not os.path.isfile(args.video):
|
|
print(f"Error: video file not found: {args.video}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
selected: list[ModelConfig] = []
|
|
if args.all_models:
|
|
selected = list(discovered)
|
|
if args.model:
|
|
for name in args.model:
|
|
exact = [c for c in discovered if c.filename == name]
|
|
if exact:
|
|
for cfg in exact:
|
|
if cfg not in selected:
|
|
selected.append(cfg)
|
|
continue
|
|
partial = [c for c in discovered if name in c.filename]
|
|
if partial:
|
|
for cfg in partial:
|
|
if cfg not in selected:
|
|
selected.append(cfg)
|
|
continue
|
|
print(f"Warning: model '{name}' not found in {args.models_dir}",
|
|
file=sys.stderr)
|
|
|
|
if not args.all_models and not args.model:
|
|
print("Error: specify --model, --all-models, or --list-models", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
if not selected:
|
|
print("Error: no valid models selected", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
class_filter = args.filter if args.filter else None
|
|
|
|
for model_cfg in selected:
|
|
if args.output and len(selected) == 1:
|
|
out_path = args.output
|
|
else:
|
|
out_path = os.path.join(args.output_dir, f"{model_cfg.stem}_annotated.mp4")
|
|
parent = os.path.dirname(out_path)
|
|
if parent:
|
|
os.makedirs(parent, exist_ok=True)
|
|
|
|
print(f"Processing {args.video} with model {model_cfg.filename} ...")
|
|
|
|
def progress_callback(frame_idx, total_frames, _model=model_cfg.filename):
|
|
if total_frames:
|
|
print(f"\r [{_model}] frame {frame_idx}/{total_frames}",
|
|
end="", flush=True)
|
|
else:
|
|
print(f"\r [{_model}] frame {frame_idx}", end="", flush=True)
|
|
|
|
result = run_pipeline(
|
|
video_path=args.video,
|
|
model_config=model_cfg,
|
|
output_path=out_path,
|
|
class_filter=class_filter,
|
|
sack_conf=args.sack_conf,
|
|
truck_conf=args.truck_conf,
|
|
output_max_height=args.output_height,
|
|
progress_callback=progress_callback,
|
|
)
|
|
print()
|
|
print(f"Output: {result.output_path}")
|
|
print(f"Frames: {result.frame_count}")
|
|
print(f"Loading: {result.loading_count}, "
|
|
f"Unloading: {result.unloading_count}, Net: {result.net_count}")
|
|
print(f"Batches: {result.batch_count}")
|
|
print(f"Duration: {result.duration_seconds:.1f}s")
|
|
|
|
print(f"\nDone. {len(selected)} model(s) processed.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|