diff --git a/README.md b/README.md index bb3b846..edabfa5 100644 --- a/README.md +++ b/README.md @@ -73,6 +73,20 @@ python -m river_annotation_tool.annotation_script --clip left_20230615T120000 All settings live in `config/config.yaml`. Copy `config/config.example.yaml` to get started. ```yaml +filenames: + video_in_zip: left.mp4 # video filename inside each ZIP archive + video_tmp_suffix: .mp4 # suffix for the extraction temp file + zip_extension: .zip # extension used when resolving clip names + mask: mask.png # saved water mask + metadata: metadata.json # saved survey answers + frame: frame.png # middle frame snapshot + overlay: overlay.png # frame with mask blended in green + mask_vis: mask_vis.png # greyscale mask PNG (--extras only) + gif_original_hires: video_original_hires.gif + gif_original_lowres: video_original_lowres.gif + gif_overlay_hires: video_overlay_hires.gif + gif_overlay_lowres: video_overlay_lowres.gif + display_max: 720 # longest side in pixels for display fps_fallback: 25 # FPS to use if the video header is missing max_frames: 100 # max frames to extract per clip @@ -139,6 +153,8 @@ video_overlay_hires.gif # Overlay GIF at display resolution video_overlay_lowres.gif # Overlay GIF at 50% of display resolution ``` +All output filenames can be overridden via the `filenames:` section in `config/config.yaml`. + ### Survey answers (`metadata.json`) Keys and values are determined by the `questions` section in `config/config.yaml`. With the default config: @@ -159,7 +175,7 @@ Keys and values are determined by the `questions` section in `config/config.yaml ### Clip format -Each clip is a ZIP archive containing a `left.mp4` video. The filename encodes the recording timestamp (e.g. `left_20230615T120000.zip`). +Each clip is a ZIP archive containing a video file (default `left.mp4`, configurable via `filenames.video_in_zip`). The filename encodes the recording timestamp (e.g. `left_20230615T120000.zip`). ### Frame loading diff --git a/config/config.example.yaml b/config/config.example.yaml index b34a59e..18c3c4c 100644 --- a/config/config.example.yaml +++ b/config/config.example.yaml @@ -1,3 +1,17 @@ +filenames: + video_in_zip: left.mp4 + video_tmp_suffix: .mp4 + zip_extension: .zip + mask: mask.png + metadata: metadata.json + frame: frame.png + overlay: overlay.png + mask_vis: mask_vis.png + gif_original_hires: video_original_hires.gif + gif_original_lowres: video_original_lowres.gif + gif_overlay_hires: video_overlay_hires.gif + gif_overlay_lowres: video_overlay_lowres.gif + display_max: 720 fps_fallback: 25 max_frames: 100 diff --git a/src/river_annotation_tool/annotation_script.py b/src/river_annotation_tool/annotation_script.py index 4d43146..72b3ab7 100644 --- a/src/river_annotation_tool/annotation_script.py +++ b/src/river_annotation_tool/annotation_script.py @@ -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", diff --git a/src/river_annotation_tool/annotator.py b/src/river_annotation_tool/annotator.py index cc3f6a6..921920e 100644 --- a/src/river_annotation_tool/annotator.py +++ b/src/river_annotation_tool/annotator.py @@ -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) diff --git a/src/river_annotation_tool/clip_selector.py b/src/river_annotation_tool/clip_selector.py index 775d593..c7ef319 100644 --- a/src/river_annotation_tool/clip_selector.py +++ b/src/river_annotation_tool/clip_selector.py @@ -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 [] diff --git a/src/river_annotation_tool/config.py b/src/river_annotation_tool/config.py index dddcc28..46f1c59 100644 --- a/src/river_annotation_tool/config.py +++ b/src/river_annotation_tool/config.py @@ -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 diff --git a/src/river_annotation_tool/video_loader.py b/src/river_annotation_tool/video_loader.py index 75b437e..59a07dc 100644 --- a/src/river_annotation_tool/video_loader.py +++ b/src/river_annotation_tool/video_loader.py @@ -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