613 lines
19 KiB
Python
613 lines
19 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""
|
|||
|
|
从视频读取连续帧，使用 OpenCV Farneback 稠密光流估计运动，
|
|||
|
|
输出带光流伪彩色（方向=色调、速度=亮度）的可视化视频。
|
|||
|
|
|
|||
|
|
声呐类画面背景噪声也会在光流里产生响应；可用 --fish-mask 仅在高亮「鱼团」
|
|||
|
|
区域显示光流，并用连通域面积过滤掉细碎亮点。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
import argparse
|
|||
|
|
import logging
|
|||
|
|
import time
|
|||
|
|
from pathlib import Path
|
|||
|
|
from types import SimpleNamespace
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
|
|||
|
|
_log = logging.getLogger("optical_flow")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def flow_to_bgr(
|
|||
|
|
flow: np.ndarray,
|
|||
|
|
mag_clip_percentile: float = 95.0,
|
|||
|
|
valid_mask: np.ndarray | None = None,
|
|||
|
|
) -> np.ndarray:
|
|||
|
|
"""将 (H,W,2) 光流转为 BGR 伪彩色图。若给定 valid_mask，幅值分位数仅在掩膜内统计。"""
|
|||
|
|
fx = flow[..., 0].astype(np.float32)
|
|||
|
|
fy = flow[..., 1].astype(np.float32)
|
|||
|
|
mag = np.sqrt(fx * fx + fy * fy)
|
|||
|
|
ang = np.arctan2(fy, fx)
|
|||
|
|
|
|||
|
|
if valid_mask is not None and valid_mask.size > 0:
|
|||
|
|
vm = valid_mask > 0
|
|||
|
|
if np.any(vm):
|
|||
|
|
clip = float(np.percentile(mag[vm], mag_clip_percentile))
|
|||
|
|
else:
|
|||
|
|
clip = float(np.percentile(mag, mag_clip_percentile))
|
|||
|
|
else:
|
|||
|
|
clip = float(np.percentile(mag, mag_clip_percentile))
|
|||
|
|
if clip < 1e-6:
|
|||
|
|
clip = 1e-6
|
|||
|
|
mag_norm = np.clip(mag / clip, 0.0, 1.0)
|
|||
|
|
|
|||
|
|
h, w = flow.shape[:2]
|
|||
|
|
hsv = np.zeros((h, w, 3), dtype=np.float32)
|
|||
|
|
hsv[..., 0] = (ang + np.pi) / (2.0 * np.pi) * 179.0
|
|||
|
|
hsv[..., 1] = 255.0
|
|||
|
|
hsv[..., 2] = mag_norm * 255.0
|
|||
|
|
hsv_u8 = hsv.astype(np.uint8)
|
|||
|
|
return cv2.cvtColor(hsv_u8, cv2.COLOR_HSV2BGR)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _odd_k(k: int) -> int:
|
|||
|
|
k = max(1, int(k))
|
|||
|
|
return k if k % 2 == 1 else k + 1
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_fish_mask(
|
|||
|
|
gray_u8: np.ndarray,
|
|||
|
|
*,
|
|||
|
|
bright_percentile: float,
|
|||
|
|
min_blob_area: int,
|
|||
|
|
open_k: int,
|
|||
|
|
close_k: int,
|
|||
|
|
dilate_k: int,
|
|||
|
|
blur_sigma: float,
|
|||
|
|
keep_largest_blobs: int = 0,
|
|||
|
|
) -> np.ndarray:
|
|||
|
|
"""
|
|||
|
|
声呐/暗背景高亮目标：取灰度高分位作为阈值，形态学去噪，保留足够大的连通域。
|
|||
|
|
"""
|
|||
|
|
g = gray_u8
|
|||
|
|
if blur_sigma > 1e-6:
|
|||
|
|
k = _odd_k(int(round(blur_sigma * 6)) | 1)
|
|||
|
|
k = max(3, min(k, 31))
|
|||
|
|
g = cv2.GaussianBlur(g, (k, k), blur_sigma)
|
|||
|
|
|
|||
|
|
thr = float(np.percentile(g.astype(np.float32), bright_percentile))
|
|||
|
|
binary = (g.astype(np.float32) >= thr).astype(np.uint8) * 255
|
|||
|
|
|
|||
|
|
if open_k > 0:
|
|||
|
|
ok = _odd_k(open_k)
|
|||
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (ok, ok))
|
|||
|
|
binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)
|
|||
|
|
if close_k > 0:
|
|||
|
|
ck = _odd_k(close_k)
|
|||
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (ck, ck))
|
|||
|
|
binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
|
|||
|
|
|
|||
|
|
num, labels, stats, _ = cv2.connectedComponentsWithStats(binary, connectivity=8)
|
|||
|
|
out = np.zeros_like(binary)
|
|||
|
|
for i in range(1, num):
|
|||
|
|
if stats[i, cv2.CC_STAT_AREA] >= min_blob_area:
|
|||
|
|
out[labels == i] = 255
|
|||
|
|
|
|||
|
|
if dilate_k > 0:
|
|||
|
|
dk = _odd_k(dilate_k)
|
|||
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dk, dk))
|
|||
|
|
out = cv2.dilate(out, kernel)
|
|||
|
|
|
|||
|
|
if keep_largest_blobs > 0:
|
|||
|
|
num, labels, stats, _ = cv2.connectedComponentsWithStats(out, connectivity=8)
|
|||
|
|
ranked: list[tuple[int, int]] = []
|
|||
|
|
for i in range(1, num):
|
|||
|
|
ranked.append((int(stats[i, cv2.CC_STAT_AREA]), i))
|
|||
|
|
ranked.sort(reverse=True)
|
|||
|
|
trimmed = np.zeros_like(out)
|
|||
|
|
for _, comp_idx in ranked[:keep_largest_blobs]:
|
|||
|
|
trimmed[labels == comp_idx] = 255
|
|||
|
|
out = trimmed
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
|
|||
|
|
def apply_flow_mask(flow_bgr: np.ndarray, mask_u8: np.ndarray) -> np.ndarray:
|
|||
|
|
"""掩膜外置黑；三通道与 mask 相乘。"""
|
|||
|
|
m = (mask_u8.astype(np.float32) / 255.0)[..., np.newaxis]
|
|||
|
|
return np.clip(flow_bgr.astype(np.float32) * m, 0, 255).astype(np.uint8)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def draw_flow_arrows(
|
|||
|
|
flow: np.ndarray,
|
|||
|
|
canvas_bgr: np.ndarray,
|
|||
|
|
*,
|
|||
|
|
mask_u8: np.ndarray | None,
|
|||
|
|
step: int,
|
|||
|
|
scale: float,
|
|||
|
|
min_magnitude: float,
|
|||
|
|
color: tuple[int, int, int],
|
|||
|
|
thickness: int,
|
|||
|
|
) -> np.ndarray:
|
|||
|
|
"""
|
|||
|
|
将稠密光流按网格采样绘制为箭头，便于精确观察局部方向与速度。
|
|||
|
|
"""
|
|||
|
|
out = canvas_bgr.copy()
|
|||
|
|
h, w = flow.shape[:2]
|
|||
|
|
s = max(2, int(step))
|
|||
|
|
t = max(1, int(thickness))
|
|||
|
|
mag2_thr = float(min_magnitude) * float(min_magnitude)
|
|||
|
|
|
|||
|
|
for y in range(s // 2, h, s):
|
|||
|
|
for x in range(s // 2, w, s):
|
|||
|
|
if mask_u8 is not None and mask_u8[y, x] == 0:
|
|||
|
|
continue
|
|||
|
|
fx = float(flow[y, x, 0])
|
|||
|
|
fy = float(flow[y, x, 1])
|
|||
|
|
if fx * fx + fy * fy < mag2_thr:
|
|||
|
|
continue
|
|||
|
|
x2 = int(round(x + fx * scale))
|
|||
|
|
y2 = int(round(y + fy * scale))
|
|||
|
|
x2 = max(0, min(w - 1, x2))
|
|||
|
|
y2 = max(0, min(h - 1, y2))
|
|||
|
|
cv2.arrowedLine(
|
|||
|
|
out,
|
|||
|
|
(x, y),
|
|||
|
|
(x2, y2),
|
|||
|
|
color=color,
|
|||
|
|
thickness=t,
|
|||
|
|
tipLength=0.3,
|
|||
|
|
)
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_output_frame(
|
|||
|
|
frame_bgr: np.ndarray,
|
|||
|
|
flow_bgr: np.ndarray,
|
|||
|
|
mode: str,
|
|||
|
|
) -> np.ndarray:
|
|||
|
|
if mode == "sidebyside":
|
|||
|
|
return np.hstack([frame_bgr, flow_bgr])
|
|||
|
|
if mode == "overlay":
|
|||
|
|
return cv2.addWeighted(frame_bgr, 0.55, flow_bgr, 0.45, 0.0)
|
|||
|
|
raise ValueError(f"未知 mode: {mode}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _try_gst_nvenc_writer(
|
|||
|
|
path: Path,
|
|||
|
|
fps: float,
|
|||
|
|
size: tuple[int, int],
|
|||
|
|
) -> cv2.VideoWriter | None:
|
|||
|
|
"""Try to create a cv2.VideoWriter backed by GStreamer + Jetson NVENC.
|
|||
|
|
|
|||
|
|
Returns None if GStreamer or the hardware encoder is unavailable.
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
if not hasattr(cv2, "CAP_GSTREAMER"):
|
|||
|
|
return None
|
|||
|
|
w, h = size
|
|||
|
|
loc = str(path).replace('"', '\\"')
|
|||
|
|
gst_pipe = (
|
|||
|
|
f'appsrc ! videoconvert ! video/x-raw,format=BGRx ! '
|
|||
|
|
f'nvvidconv ! video/x-raw(memory:NVMM) ! '
|
|||
|
|
f'nvv4l2h264enc bitrate=4000000 ! h264parse ! '
|
|||
|
|
f'mp4mux ! filesink location="{loc}"'
|
|||
|
|
)
|
|||
|
|
writer = cv2.VideoWriter(gst_pipe, cv2.CAP_GSTREAMER, 0, fps, (w, h))
|
|||
|
|
if writer.isOpened():
|
|||
|
|
_log.info("[optical-flow] using GStreamer NVENC writer for %s", path.name)
|
|||
|
|
return writer
|
|||
|
|
writer.release()
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def open_writer(
|
|||
|
|
path: Path,
|
|||
|
|
fps: float,
|
|||
|
|
size: tuple[int, int],
|
|||
|
|
fourcc_str: str,
|
|||
|
|
) -> cv2.VideoWriter:
|
|||
|
|
gst = _try_gst_nvenc_writer(path, fps, size)
|
|||
|
|
if gst is not None:
|
|||
|
|
return gst
|
|||
|
|
fourcc = cv2.VideoWriter_fourcc(*fourcc_str)
|
|||
|
|
writer = cv2.VideoWriter(str(path), fourcc, fps, size)
|
|||
|
|
if not writer.isOpened():
|
|||
|
|
raise RuntimeError(
|
|||
|
|
f"无法创建输出视频: {path}（codec={fourcc_str}）。"
|
|||
|
|
"可尝试 --fourcc avc1 或 XVID 并配合 .avi 扩展名。"
|
|||
|
|
)
|
|||
|
|
return writer
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _validate_flow_args(args: Any) -> None:
|
|||
|
|
if not isinstance(args.resize, (int, float)) or float(args.resize) <= 0:
|
|||
|
|
raise ValueError("--resize 必须为正数")
|
|||
|
|
if not 0 < float(args.bright_percentile) < 100:
|
|||
|
|
raise ValueError("--bright-percentile 应在 (0, 100) 内")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _run_flow_core(args: Any) -> None:
|
|||
|
|
"""稠密光流 + 可视化；``args`` 需含与 CLI 相同的字段（含 ``input`` / ``output`` 路径）。"""
|
|||
|
|
in_path = Path(args.input)
|
|||
|
|
out_path = Path(args.output)
|
|||
|
|
if not in_path.is_file():
|
|||
|
|
raise FileNotFoundError(f"找不到输入视频: {in_path}")
|
|||
|
|
|
|||
|
|
src_size_mb = in_path.stat().st_size / (1024 * 1024)
|
|||
|
|
_log.info("[optical-flow] start: %s (%.1f MB), mode=%s, resize=%.2f",
|
|||
|
|
in_path.name, src_size_mb, args.mode, args.resize)
|
|||
|
|
t0 = time.monotonic()
|
|||
|
|
|
|||
|
|
cap = cv2.VideoCapture(str(in_path))
|
|||
|
|
if not cap.isOpened():
|
|||
|
|
raise RuntimeError(f"无法打开视频: {in_path}")
|
|||
|
|
|
|||
|
|
writer: cv2.VideoWriter | None = None
|
|||
|
|
try:
|
|||
|
|
fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
|||
|
|
w_in = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|||
|
|
h_in = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|||
|
|
_log.info("[optical-flow] src %dx%d @ %.1f fps", w_in, h_in, fps)
|
|||
|
|
scale = float(args.resize)
|
|||
|
|
|
|||
|
|
w = max(1, int(round(w_in * scale)))
|
|||
|
|
h = max(1, int(round(h_in * scale)))
|
|||
|
|
|
|||
|
|
if args.mode == "sidebyside":
|
|||
|
|
out_w, out_h = w * 2, h
|
|||
|
|
else:
|
|||
|
|
out_w, out_h = w, h
|
|||
|
|
|
|||
|
|
writer = open_writer(out_path, fps, (out_w, out_h), args.fourcc)
|
|||
|
|
|
|||
|
|
ret, prev_bgr = cap.read()
|
|||
|
|
if not ret:
|
|||
|
|
raise RuntimeError("视频为空或无法读取首帧")
|
|||
|
|
|
|||
|
|
if scale != 1.0:
|
|||
|
|
prev_bgr = cv2.resize(prev_bgr, (w, h), interpolation=cv2.INTER_AREA)
|
|||
|
|
prev_gray = cv2.cvtColor(prev_bgr, cv2.COLOR_BGR2GRAY)
|
|||
|
|
|
|||
|
|
black_flow = np.zeros((h, w, 3), dtype=np.uint8)
|
|||
|
|
first_out = build_output_frame(prev_bgr, black_flow, args.mode)
|
|||
|
|
writer.write(first_out)
|
|||
|
|
written = 1
|
|||
|
|
frame_idx = 0
|
|||
|
|
|
|||
|
|
fb_flags = 0
|
|||
|
|
progress_log = bool(getattr(args, "progress_log", False))
|
|||
|
|
while True:
|
|||
|
|
ret, frame_bgr = cap.read()
|
|||
|
|
if not ret:
|
|||
|
|
break
|
|||
|
|
if scale != 1.0:
|
|||
|
|
frame_bgr = cv2.resize(frame_bgr, (w, h), interpolation=cv2.INTER_AREA)
|
|||
|
|
gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)
|
|||
|
|
|
|||
|
|
flow = cv2.calcOpticalFlowFarneback(
|
|||
|
|
prev_gray,
|
|||
|
|
gray,
|
|||
|
|
None,
|
|||
|
|
args.pyr_scale,
|
|||
|
|
args.levels,
|
|||
|
|
args.winsize,
|
|||
|
|
args.iterations,
|
|||
|
|
args.poly_n,
|
|||
|
|
args.poly_sigma,
|
|||
|
|
fb_flags,
|
|||
|
|
)
|
|||
|
|
mask: np.ndarray | None = None
|
|||
|
|
if args.fish_mask:
|
|||
|
|
mask = build_fish_mask(
|
|||
|
|
gray,
|
|||
|
|
bright_percentile=args.bright_percentile,
|
|||
|
|
min_blob_area=args.min_blob_area,
|
|||
|
|
open_k=args.mask_open,
|
|||
|
|
close_k=args.mask_close,
|
|||
|
|
dilate_k=args.mask_dilate,
|
|||
|
|
blur_sigma=args.mask_blur,
|
|||
|
|
keep_largest_blobs=args.keep_largest_blobs,
|
|||
|
|
)
|
|||
|
|
has_mask = np.count_nonzero(mask) > 0
|
|||
|
|
else:
|
|||
|
|
has_mask = False
|
|||
|
|
|
|||
|
|
if args.viz_style in ("hsv", "hsv_arrows"):
|
|||
|
|
if args.fish_mask and has_mask and mask is not None:
|
|||
|
|
hsv_bgr = flow_to_bgr(flow, valid_mask=mask)
|
|||
|
|
hsv_bgr = apply_flow_mask(hsv_bgr, mask)
|
|||
|
|
elif args.fish_mask:
|
|||
|
|
hsv_bgr = np.zeros((h, w, 3), dtype=np.uint8)
|
|||
|
|
else:
|
|||
|
|
hsv_bgr = flow_to_bgr(flow)
|
|||
|
|
else:
|
|||
|
|
hsv_bgr = np.zeros((h, w, 3), dtype=np.uint8)
|
|||
|
|
|
|||
|
|
if args.viz_style in ("arrows", "hsv_arrows"):
|
|||
|
|
arrow_base = np.zeros((h, w, 3), dtype=np.uint8)
|
|||
|
|
arrow_bgr = draw_flow_arrows(
|
|||
|
|
flow,
|
|||
|
|
arrow_base,
|
|||
|
|
mask_u8=mask if (args.fish_mask and has_mask) else None,
|
|||
|
|
step=args.arrow_step,
|
|||
|
|
scale=args.arrow_scale,
|
|||
|
|
min_magnitude=args.arrow_threshold,
|
|||
|
|
color=(0, 255, 255),
|
|||
|
|
thickness=args.arrow_thickness,
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
arrow_bgr = np.zeros((h, w, 3), dtype=np.uint8)
|
|||
|
|
|
|||
|
|
flow_bgr = cv2.addWeighted(hsv_bgr, 0.8, arrow_bgr, 1.0, 0.0)
|
|||
|
|
out_frame = build_output_frame(frame_bgr, flow_bgr, args.mode)
|
|||
|
|
writer.write(out_frame)
|
|||
|
|
written += 1
|
|||
|
|
|
|||
|
|
prev_gray = gray
|
|||
|
|
frame_idx += 1
|
|||
|
|
if progress_log and frame_idx % 30 == 0:
|
|||
|
|
print(f"已处理 {frame_idx} 帧光流…", flush=True)
|
|||
|
|
|
|||
|
|
elapsed = time.monotonic() - t0
|
|||
|
|
_log.info("[optical-flow] done: %d frames in %.1fs (%.1f fps) -> %s",
|
|||
|
|
written, elapsed, written / max(elapsed, 0.001), out_path.name)
|
|||
|
|
if progress_log:
|
|||
|
|
print(f"完成，共写入 {written} 帧（{elapsed:.1f}s），保存至: {out_path}", flush=True)
|
|||
|
|
finally:
|
|||
|
|
cap.release()
|
|||
|
|
if writer is not None:
|
|||
|
|
writer.release()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def run_optical_flow_video(
|
|||
|
|
input_path: Path,
|
|||
|
|
output_path: Path,
|
|||
|
|
*,
|
|||
|
|
mode: str = "overlay",
|
|||
|
|
viz_style: str = "hsv",
|
|||
|
|
resize: float = 1.0,
|
|||
|
|
fourcc: str = "mp4v",
|
|||
|
|
pyr_scale: float = 0.5,
|
|||
|
|
levels: int = 3,
|
|||
|
|
winsize: int = 15,
|
|||
|
|
iterations: int = 3,
|
|||
|
|
poly_n: int = 5,
|
|||
|
|
poly_sigma: float = 1.2,
|
|||
|
|
fish_mask: bool = True,
|
|||
|
|
bright_percentile: float = 97.5,
|
|||
|
|
min_blob_area: int = 500,
|
|||
|
|
mask_open: int = 3,
|
|||
|
|
mask_close: int = 11,
|
|||
|
|
mask_dilate: int = 5,
|
|||
|
|
mask_blur: float = 1.0,
|
|||
|
|
keep_largest_blobs: int = 0,
|
|||
|
|
arrow_step: int = 12,
|
|||
|
|
arrow_scale: float = 2.0,
|
|||
|
|
arrow_threshold: float = 0.8,
|
|||
|
|
arrow_thickness: int = 1,
|
|||
|
|
progress_log: bool = False,
|
|||
|
|
) -> bool:
|
|||
|
|
"""供 fish_api 等调用：成功写出 ``output_path`` 返回 True，否则 False。"""
|
|||
|
|
args = SimpleNamespace(
|
|||
|
|
input=input_path,
|
|||
|
|
output=output_path,
|
|||
|
|
mode=mode,
|
|||
|
|
viz_style=viz_style,
|
|||
|
|
resize=resize,
|
|||
|
|
fourcc=fourcc,
|
|||
|
|
pyr_scale=pyr_scale,
|
|||
|
|
levels=levels,
|
|||
|
|
winsize=winsize,
|
|||
|
|
iterations=iterations,
|
|||
|
|
poly_n=poly_n,
|
|||
|
|
poly_sigma=poly_sigma,
|
|||
|
|
fish_mask=fish_mask,
|
|||
|
|
bright_percentile=bright_percentile,
|
|||
|
|
min_blob_area=min_blob_area,
|
|||
|
|
mask_open=mask_open,
|
|||
|
|
mask_close=mask_close,
|
|||
|
|
mask_dilate=mask_dilate,
|
|||
|
|
mask_blur=mask_blur,
|
|||
|
|
keep_largest_blobs=keep_largest_blobs,
|
|||
|
|
arrow_step=arrow_step,
|
|||
|
|
arrow_scale=arrow_scale,
|
|||
|
|
arrow_threshold=arrow_threshold,
|
|||
|
|
arrow_thickness=arrow_thickness,
|
|||
|
|
progress_log=progress_log,
|
|||
|
|
)
|
|||
|
|
try:
|
|||
|
|
_validate_flow_args(args)
|
|||
|
|
_run_flow_core(args)
|
|||
|
|
except Exception:
|
|||
|
|
_log.exception("[optical-flow] run_optical_flow_video failed: %s -> %s",
|
|||
|
|
input_path, output_path)
|
|||
|
|
return False
|
|||
|
|
return output_path.is_file() and output_path.stat().st_size > 0
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main() -> None:
|
|||
|
|
script_dir = Path(__file__).resolve().parent
|
|||
|
|
default_in = script_dir / "fish_echo.MP4"
|
|||
|
|
default_out = script_dir / "fish_echo_flow_vis.mp4"
|
|||
|
|
|
|||
|
|
p = argparse.ArgumentParser(description="视频稠密光流可视化（OpenCV Farneback）")
|
|||
|
|
p.add_argument(
|
|||
|
|
"--input",
|
|||
|
|
"-i",
|
|||
|
|
type=Path,
|
|||
|
|
default=default_in,
|
|||
|
|
help=f"输入视频路径（默认: {default_in.name}）",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--output",
|
|||
|
|
"-o",
|
|||
|
|
type=Path,
|
|||
|
|
default=default_out,
|
|||
|
|
help=f"输出视频路径（默认: {default_out.name}）",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--mode",
|
|||
|
|
choices=("sidebyside", "overlay"),
|
|||
|
|
default="sidebyside",
|
|||
|
|
help="sidebyside: 原图|光流；overlay: 原图与光流半透明叠加",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--viz-style",
|
|||
|
|
choices=("hsv", "arrows", "hsv_arrows"),
|
|||
|
|
default="hsv",
|
|||
|
|
help="光流可视化风格：hsv 伪彩色、arrows 箭头、hsv_arrows 组合",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--resize",
|
|||
|
|
type=float,
|
|||
|
|
default=1.0,
|
|||
|
|
metavar="SCALE",
|
|||
|
|
help="处理前将帧宽高乘以该比例以加速（例如 0.5）",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--fourcc",
|
|||
|
|
default="mp4v",
|
|||
|
|
help="VideoWriter 四字符编码，常见: mp4v, avc1, XVID",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--pyr-scale",
|
|||
|
|
type=float,
|
|||
|
|
default=0.5,
|
|||
|
|
help="Farneback 金字塔缩放（OpenCV pyr_scale）",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--levels",
|
|||
|
|
type=int,
|
|||
|
|
default=3,
|
|||
|
|
help="Farneback 金字塔层数",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--winsize",
|
|||
|
|
type=int,
|
|||
|
|
default=15,
|
|||
|
|
help="Farneback 窗口大小",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--iterations",
|
|||
|
|
type=int,
|
|||
|
|
default=3,
|
|||
|
|
help="Farneback 每层迭代次数",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--poly-n",
|
|||
|
|
type=int,
|
|||
|
|
default=5,
|
|||
|
|
help="Farneback 像素邻域大小（poly_n）",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--poly-sigma",
|
|||
|
|
type=float,
|
|||
|
|
default=1.2,
|
|||
|
|
help="Farneback 高斯标准差（poly_sigma）",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--fish-mask",
|
|||
|
|
dest="fish_mask",
|
|||
|
|
action="store_true",
|
|||
|
|
default=True,
|
|||
|
|
help="仅在高亮鱼团区域显示光流（默认开启，抑制背景噪声）",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--no-fish-mask",
|
|||
|
|
dest="fish_mask",
|
|||
|
|
action="store_false",
|
|||
|
|
help="关闭鱼团掩膜，整幅画面显示光流（与旧行为一致）",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--bright-percentile",
|
|||
|
|
type=float,
|
|||
|
|
default=97.5,
|
|||
|
|
metavar="P",
|
|||
|
|
help="灰度阈值：取 >= 该分位数的像素作为「亮斑」候选（越高越只保留最亮区域）",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--min-blob-area",
|
|||
|
|
type=int,
|
|||
|
|
default=500,
|
|||
|
|
metavar="PX",
|
|||
|
|
help="连通域最小面积（像素），小于此面积的亮斑视为噪声并丢弃",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--mask-open",
|
|||
|
|
type=int,
|
|||
|
|
default=3,
|
|||
|
|
help="形态学开运算核大小（奇数，0 表示跳过），去小噪点",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--mask-close",
|
|||
|
|
type=int,
|
|||
|
|
default=11,
|
|||
|
|
help="形态学闭运算核大小（奇数，0 表示跳过），连接同一鱼团碎片",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--mask-dilate",
|
|||
|
|
type=int,
|
|||
|
|
default=5,
|
|||
|
|
help="掩膜膨胀核大小（奇数，0 表示不膨胀），略扩大显示区域包住边缘运动",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--mask-blur",
|
|||
|
|
type=float,
|
|||
|
|
default=1.0,
|
|||
|
|
help="阈值前高斯模糊 sigma（0 表示不模糊），可平滑细碎纹理",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--keep-largest-blobs",
|
|||
|
|
type=int,
|
|||
|
|
default=0,
|
|||
|
|
metavar="N",
|
|||
|
|
help="在面积过滤后仅保留面积最大的 N 个连通域（0 表示不限制，适合多鱼场景）",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--arrow-step",
|
|||
|
|
type=int,
|
|||
|
|
default=12,
|
|||
|
|
help="箭头网格步长（像素），越小越密",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--arrow-scale",
|
|||
|
|
type=float,
|
|||
|
|
default=2.0,
|
|||
|
|
help="箭头长度缩放系数",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--arrow-threshold",
|
|||
|
|
type=float,
|
|||
|
|
default=0.8,
|
|||
|
|
help="绘制箭头的最小速度阈值（像素/帧）",
|
|||
|
|
)
|
|||
|
|
p.add_argument(
|
|||
|
|
"--arrow-thickness",
|
|||
|
|
type=int,
|
|||
|
|
default=1,
|
|||
|
|
help="箭头线宽",
|
|||
|
|
)
|
|||
|
|
args = p.parse_args()
|
|||
|
|
if not args.input.is_file():
|
|||
|
|
raise SystemExit(f"找不到输入视频: {args.input}")
|
|||
|
|
try:
|
|||
|
|
_validate_flow_args(args)
|
|||
|
|
except ValueError as e:
|
|||
|
|
raise SystemExit(str(e)) from e
|
|||
|
|
args.progress_log = True
|
|||
|
|
try:
|
|||
|
|
_run_flow_core(args)
|
|||
|
|
except FileNotFoundError as e:
|
|||
|
|
raise SystemExit(str(e)) from e
|
|||
|
|
except RuntimeError as e:
|
|||
|
|
raise SystemExit(str(e)) from e
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|