All constants are in config

This commit is contained in:
2026-05-20 14:00:11 +02:00
parent 6a0259c6cf
commit b4daa28354
7 changed files with 96 additions and 23 deletions

View File

@@ -23,9 +23,7 @@ def parse_args():
parser.add_argument("--out", default=None, help="Override out_dir from config")
parser.add_argument("--clips", default=None, help="Override clips_file from config")
parser.add_argument(
"--clip",
default=None,
help="Stem name of a specific clip to load (e.g. 'left_20230501')",
"--clip", default=None, help="Stem name of a specific clip to load"
)
parser.add_argument(
"--extras",

View File

@@ -40,6 +40,8 @@ class Annotator(QMainWindow):
data_dir=Path(config.data_dir),
out_dir=self.out_dir,
clips_file=Path(config.clips_file),
mask_filename=config.filenames.mask,
zip_extension=config.filenames.zip_extension,
)
self.setWindowTitle("River Annotator")
@@ -55,11 +57,13 @@ class Annotator(QMainWindow):
self.cfg.max_frames,
self.cfg.display_max,
self.cfg.fps_fallback,
self.cfg.filenames.video_in_zip,
self.cfg.filenames.video_tmp_suffix,
)
self._pending_answers = self._read_saved_answers()
def _read_saved_mask(self):
mask_path = self.out_dir / self.filename.stem / "mask.png"
mask_path = self.out_dir / self.filename.stem / self.cfg.filenames.mask
if not mask_path.exists():
return None
mask_full = np.array(Image.open(mask_path).convert("L"))
@@ -70,7 +74,7 @@ class Annotator(QMainWindow):
)
def _read_saved_answers(self):
meta_path = self.out_dir / self.filename.stem / "metadata.json"
meta_path = self.out_dir / self.filename.stem / self.cfg.filenames.metadata
if not meta_path.exists():
return None
with open(meta_path) as f:
@@ -214,23 +218,26 @@ class Annotator(QMainWindow):
(self.w, self.h),
interpolation=cv2.INTER_NEAREST,
)
Image.fromarray(mask_full * 255).save(out / "mask.png")
fn = self.cfg.filenames
Image.fromarray(mask_full * 255).save(out / fn.mask)
with open(out / "metadata.json", "w") as f:
with open(out / fn.metadata, "w") as f:
json.dump(self.get_answers(), f, indent=2)
mid = len(self.frames) // 2
frame = self.frames[mid]
Image.fromarray(frame).save(out / "frame.png")
Image.fromarray(self._make_overlay(frame)).save(out / "overlay.png")
Image.fromarray(frame).save(out / fn.frame)
Image.fromarray(self._make_overlay(frame)).save(out / fn.overlay)
if self.extras:
Image.fromarray((self.mc.mask * 255).astype(np.uint8)).save(out / "mask_vis.png")
Image.fromarray((self.mc.mask * 255).astype(np.uint8)).save(
out / fn.mask_vis
)
overlay_frames = [self._make_overlay(f) for f in self.frames]
self._save_gif(self.frames, out / "video_original_hires.gif", scale=1.0)
self._save_gif(self.frames, out / "video_original_lowres.gif", scale=0.5)
self._save_gif(overlay_frames, out / "video_overlay_hires.gif", scale=1.0)
self._save_gif(overlay_frames, out / "video_overlay_lowres.gif", scale=0.5)
self._save_gif(self.frames, out / fn.gif_original_hires, scale=1.0)
self._save_gif(self.frames, out / fn.gif_original_lowres, scale=0.5)
self._save_gif(overlay_frames, out / fn.gif_overlay_hires, scale=1.0)
self._save_gif(overlay_frames, out / fn.gif_overlay_lowres, scale=0.5)
print("Saved:", out)

View File

@@ -2,9 +2,18 @@ from pathlib import Path
class ClipSelector:
def __init__(self, data_dir: Path, out_dir: Path, clips_file: Path):
def __init__(
self,
data_dir: Path,
out_dir: Path,
clips_file: Path,
mask_filename: str = "mask.png",
zip_extension: str = ".zip",
):
self.data_dir = data_dir
self.out_dir = out_dir
self.mask_filename = mask_filename
self.zip_extension = zip_extension
self.clips = self._load_clips(clips_file)
self.index = 0
@@ -17,7 +26,7 @@ class ClipSelector:
]
def is_annotated(self, path: Path) -> bool:
return (self.out_dir / path.stem / "mask.png").exists()
return (self.out_dir / path.stem / self.mask_filename).exists()
def next(self, specific: str = None) -> Path:
if specific:
@@ -25,7 +34,7 @@ class ClipSelector:
return self._pick_next()
def _resolve_specific(self, specific: str) -> Path:
matches = list(self.data_dir.glob(f"{specific}.zip"))
matches = list(self.data_dir.glob(f"{specific}{self.zip_extension}"))
if not matches:
p = self.data_dir / specific
matches = [p] if p.exists() else []

View File

@@ -4,6 +4,22 @@ from pathlib import Path
import yaml
@dataclass
class FilenameConfig:
video_in_zip: str = "left.mp4"
video_tmp_suffix: str = ".mp4"
zip_extension: str = ".zip"
mask: str = "mask.png"
metadata: str = "metadata.json"
frame: str = "frame.png"
overlay: str = "overlay.png"
mask_vis: str = "mask_vis.png"
gif_original_hires: str = "video_original_hires.gif"
gif_original_lowres: str = "video_original_lowres.gif"
gif_overlay_hires: str = "video_overlay_hires.gif"
gif_overlay_lowres: str = "video_overlay_lowres.gif"
@dataclass
class AppConfig:
display_max: int = 480
@@ -13,6 +29,7 @@ class AppConfig:
out_dir: str = "data/annotation_results"
clips_file: str = "config/clips.txt"
questions: list = field(default_factory=list)
filenames: FilenameConfig = field(default_factory=FilenameConfig)
def get_questions(self):
return [
@@ -23,7 +40,9 @@ class AppConfig:
item["key"],
item["label"],
[str(o) for o in item["options"]],
str(item["default"]) if item.get("default") is not None else None,
str(item["default"])
if item.get("default") is not None
else None,
)
for item in s["items"]
],
@@ -35,4 +54,7 @@ class AppConfig:
def load_config(path: Path) -> AppConfig:
with open(path) as f:
data = yaml.safe_load(f)
return AppConfig(**data)
fn_data = data.pop("filenames", {})
cfg = AppConfig(**data)
cfg.filenames = FilenameConfig(**fn_data)
return cfg

View File

@@ -6,10 +6,17 @@ from pathlib import Path
import cv2
def load_frames(zip_path: Path, max_frames: int, display_max: int, fps_fallback: int):
video_bytes = zipfile.ZipFile(zip_path).read("left.mp4")
def load_frames(
zip_path: Path,
max_frames: int,
display_max: int,
fps_fallback: int,
video_in_zip: str = "left.mp4",
video_tmp_suffix: str = ".mp4",
):
video_bytes = zipfile.ZipFile(zip_path).read(video_in_zip)
with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f:
with tempfile.NamedTemporaryFile(suffix=video_tmp_suffix, delete=False) as f:
f.write(video_bytes)
tmp_path = f.name