Add optical flow Auto Segment button
Annotators can now press Auto Segment to replace the current mask with an automatic river segmentation based on dense optical flow magnitude and frame brightness. The result is pushed onto the undo stack, so it can be refined or reverted like any other mask operation. Parameters (norm_squared_threshold, gaussian_kernel, brightness_range) live in a separate config/optical_flow_config.yaml; the button is only enabled when optical_flow_config_file is set in config.yaml. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -20,7 +20,8 @@ from PySide6.QtWidgets import (
|
||||
)
|
||||
|
||||
from .clip_selector import ClipSelector
|
||||
from .config import AppConfig
|
||||
from .compute_optical_flow import compute_optical_flow_mask
|
||||
from .config import AppConfig, load_optical_flow_config
|
||||
from .mask_canvas import MaskCanvas
|
||||
from .video_loader import load_frames
|
||||
|
||||
@@ -38,6 +39,11 @@ class Annotator(QMainWindow):
|
||||
self.cfg = config
|
||||
self.out_dir = Path(config.out_dir)
|
||||
self.extras = extras
|
||||
self.of_cfg = (
|
||||
load_optical_flow_config(Path(config.optical_flow_config_file))
|
||||
if config.optical_flow_config_file
|
||||
else None
|
||||
)
|
||||
|
||||
self.selector = ClipSelector(
|
||||
data_dir=Path(config.data_dir),
|
||||
@@ -114,9 +120,17 @@ class Annotator(QMainWindow):
|
||||
btn_undo10 = QPushButton("Undo×10")
|
||||
btn_redo = QPushButton("Redo")
|
||||
btn_load_prev_mask = QPushButton("Load Prev Mask")
|
||||
btn_auto_segment = QPushButton("Auto Segment")
|
||||
btn_auto_segment.setEnabled(self.of_cfg is not None and self.of_cfg.enabled)
|
||||
|
||||
row1 = QHBoxLayout()
|
||||
for b in [self.btn_prev, btn_next, btn_skip, btn_load_prev_mask]:
|
||||
for b in [
|
||||
self.btn_prev,
|
||||
btn_next,
|
||||
btn_skip,
|
||||
btn_load_prev_mask,
|
||||
btn_auto_segment,
|
||||
]:
|
||||
row1.addWidget(b)
|
||||
|
||||
row_tools = QHBoxLayout()
|
||||
@@ -198,6 +212,7 @@ class Annotator(QMainWindow):
|
||||
btn_undo10.clicked.connect(self.mc.undo10)
|
||||
btn_redo.clicked.connect(self.mc.redo)
|
||||
btn_load_prev_mask.clicked.connect(self.load_prev_mask)
|
||||
btn_auto_segment.clicked.connect(self.run_optical_flow)
|
||||
|
||||
if self._pending_answers:
|
||||
self._set_answers(self._pending_answers)
|
||||
@@ -406,3 +421,13 @@ class Annotator(QMainWindow):
|
||||
interpolation=cv2.INTER_NEAREST,
|
||||
)
|
||||
self.mc.set_mask(mask)
|
||||
|
||||
def run_optical_flow(self):
|
||||
mask = compute_optical_flow_mask(
|
||||
self.frames,
|
||||
self.fps,
|
||||
self.of_cfg.norm_squared_threshold,
|
||||
self.of_cfg.gaussian_kernel,
|
||||
self.of_cfg.brightness_range,
|
||||
)
|
||||
self.mc.set_mask(mask)
|
||||
|
||||
Reference in New Issue
Block a user