Add S3 storage support via s3fs; make storage field required

- New filesystem.py: make_fs() factory (returns s3fs.S3FileSystem or None),
  plus fsjoin/fsstem/fsname path helpers
- config.py: storage field is now required ('local' or 's3'); load_config
  raises a clear ValueError when it is missing
- video_loader, clip_selector, annotator: thread fs through all file I/O;
  local paths unchanged, S3 paths use fs.open/fs.exists/fs.pipe
- annotation_script: load .env via python-dotenv at startup, create fs from
  config and pass to Annotator
- Add .env.example with SwitchEngines endpoint and AWS checksum env vars
- pyproject.toml: add s3fs and python-dotenv dependencies
- Reduce default mask alpha from 40% to 15%
- Update example clip names to colon-separated timestamps

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-05-20 16:15:38 +02:00
parent 8579bad2e2
commit dc59b8affb
15 changed files with 1539 additions and 106 deletions

View File

@@ -1,3 +1,4 @@
import io
import json
from pathlib import Path
@@ -22,6 +23,7 @@ from PySide6.QtWidgets import (
from .clip_selector import ClipSelector
from .compute_optical_flow import compute_optical_flow_mask
from .config import AppConfig, load_optical_flow_config
from .filesystem import fsjoin, fsname, fsstem
from .mask_canvas import MaskCanvas
from .video_loader import load_frames
@@ -33,11 +35,13 @@ class Annotator(QMainWindow):
clip: str = None,
extras: bool = False,
skip_annotated: bool = True,
fs=None,
):
super().__init__()
self.cfg = config
self.out_dir = Path(config.out_dir)
self.fs = fs
self.out_dir = config.out_dir
self.extras = extras
self.of_cfg = (
load_optical_flow_config(Path(config.optical_flow_config_file))
@@ -46,15 +50,16 @@ class Annotator(QMainWindow):
)
self.selector = ClipSelector(
data_dir=Path(config.data_dir),
data_dir=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,
skip_annotated=skip_annotated,
fs=fs,
)
self.history: list[Path] = []
self.history: list[str] = []
self.history_pos: int = -1
self.setWindowTitle("River Annotator")
@@ -63,8 +68,53 @@ class Annotator(QMainWindow):
self._init_ui()
self._init_timer()
# ── filesystem helpers ─────────────────────────────────────────
def _out_path(self, *parts: str) -> str:
return fsjoin(self.out_dir, fsstem(self.filename), *parts)
def _fs_exists(self, path: str) -> bool:
if self.fs is None:
return Path(path).exists()
return self.fs.exists(path)
def _fs_makedirs(self, path: str):
if self.fs is None:
Path(path).mkdir(parents=True, exist_ok=True)
else:
self.fs.makedirs(path, exist_ok=True)
def _pil_open(self, path: str) -> Image.Image:
if self.fs is None:
return Image.open(path)
with self.fs.open(path, "rb") as f:
return Image.open(io.BytesIO(f.read()))
def _pil_save(self, img: Image.Image, path: str):
if self.fs is None:
img.save(path)
else:
ext = str(path).rsplit(".", 1)[-1].lower()
fmt = "JPEG" if ext in ("jpg", "jpeg") else ext.upper()
buf = io.BytesIO()
img.save(buf, format=fmt)
self.fs.pipe(path, buf.getvalue())
def _json_read(self, path: str):
if self.fs is None:
with open(path) as f:
return json.load(f)
with self.fs.open(path, "r") as f:
return json.load(f)
def _json_write(self, data, path: str):
if self.fs is None:
with open(path, "w") as f:
json.dump(data, f, indent=2)
else:
self.fs.pipe(path, json.dumps(data, indent=2).encode())
# ── clip loading ───────────────────────────────────────────────
def _load_clip(self, specific: str = None, path: Path = None):
def _load_clip(self, specific: str = None, path: str = None):
if path is not None:
self.filename = path
else:
@@ -76,6 +126,7 @@ class Annotator(QMainWindow):
self.cfg.fps_fallback,
self.cfg.filenames.video_in_zip,
self.cfg.filenames.video_tmp_suffix,
fs=self.fs,
)
self._pending_answers = self._read_saved_answers()
@@ -85,10 +136,10 @@ class Annotator(QMainWindow):
self.history_pos = len(self.history) - 1
def _read_saved_mask(self):
mask_path = self.out_dir / self.filename.stem / self.cfg.filenames.mask
if not mask_path.exists():
mask_path = self._out_path(self.cfg.filenames.mask)
if not self._fs_exists(mask_path):
return None
mask_full = np.array(Image.open(mask_path).convert("L"))
mask_full = np.array(self._pil_open(mask_path).convert("L"))
return cv2.resize(
(mask_full > 127).astype(np.uint8),
(self.dw, self.dh),
@@ -96,16 +147,15 @@ class Annotator(QMainWindow):
)
def _read_saved_answers(self):
meta_path = self.out_dir / self.filename.stem / self.cfg.filenames.metadata
if not meta_path.exists():
meta_path = self._out_path(self.cfg.filenames.metadata)
if not self._fs_exists(meta_path):
return None
with open(meta_path) as f:
return json.load(f)
return self._json_read(meta_path)
# ── UI setup ───────────────────────────────────────────────────
def _init_ui(self):
self.mc = MaskCanvas(self.frames, self.dh, self.dw)
self.mc.set_title(self.filename.name)
self.mc.set_title(fsname(self.filename))
self.mc.reset(self._read_saved_mask())
self.q_widgets = {}
@@ -279,22 +329,34 @@ class Annotator(QMainWindow):
overlay[m] = (1 - alpha) * overlay[m] + alpha * green[m]
return overlay.astype(np.uint8)
def _save_gif(self, frames, out_path, scale=1.0):
def _save_gif(self, frames, out_path: str, scale=1.0):
h, w = frames[0].shape[:2]
nh, nw = max(1, int(h * scale)), max(1, int(w * scale))
pil_frames = [Image.fromarray(cv2.resize(f, (nw, nh))) for f in frames]
pil_frames[0].save(
out_path,
save_all=True,
append_images=pil_frames[1:],
duration=int(1000 / self.fps),
loop=0,
)
if self.fs is None:
pil_frames[0].save(
out_path,
save_all=True,
append_images=pil_frames[1:],
duration=int(1000 / self.fps),
loop=0,
)
else:
buf = io.BytesIO()
pil_frames[0].save(
buf,
format="GIF",
save_all=True,
append_images=pil_frames[1:],
duration=int(1000 / self.fps),
loop=0,
)
self.fs.pipe(out_path, buf.getvalue())
# ── actions ────────────────────────────────────────────────────
def save(self):
out = self.out_dir / self.filename.stem
out.mkdir(parents=True, exist_ok=True)
out = fsjoin(self.out_dir, fsstem(self.filename))
self._fs_makedirs(out)
mask_full = cv2.resize(
self.mc.mask.astype(np.uint8),
@@ -302,25 +364,28 @@ class Annotator(QMainWindow):
interpolation=cv2.INTER_NEAREST,
)
fn = self.cfg.filenames
Image.fromarray(mask_full * 255).save(out / fn.mask)
with open(out / fn.metadata, "w") as f:
json.dump(self.get_answers(), f, indent=2)
self._pil_save(Image.fromarray(mask_full * 255), fsjoin(out, fn.mask))
self._json_write(self.get_answers(), fsjoin(out, fn.metadata))
mid = len(self.frames) // 2
frame = self.frames[mid]
Image.fromarray(frame).save(out / fn.frame)
Image.fromarray(self._make_overlay(frame)).save(out / fn.overlay)
self._pil_save(Image.fromarray(frame), fsjoin(out, fn.frame))
self._pil_save(
Image.fromarray(self._make_overlay(frame)), fsjoin(out, fn.overlay)
)
if self.extras:
Image.fromarray((self.mc.mask * 255).astype(np.uint8)).save(
out / fn.mask_vis
self._pil_save(
Image.fromarray((self.mc.mask * 255).astype(np.uint8)),
fsjoin(out, fn.mask_vis),
)
overlay_frames = [self._make_overlay(f) for f in self.frames]
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)
self._save_gif(self.frames, fsjoin(out, fn.gif_original_hires), scale=1.0)
self._save_gif(self.frames, fsjoin(out, fn.gif_original_lowres), scale=0.5)
self._save_gif(overlay_frames, fsjoin(out, fn.gif_overlay_hires), scale=1.0)
self._save_gif(
overlay_frames, fsjoin(out, fn.gif_overlay_lowres), scale=0.5
)
print("Saved:", out)
@@ -331,7 +396,7 @@ class Annotator(QMainWindow):
self.dh,
self.dw,
mask=self._read_saved_mask(),
title=self.filename.name,
title=fsname(self.filename),
)
if self._pending_answers:
self._set_answers(self._pending_answers)
@@ -366,12 +431,12 @@ class Annotator(QMainWindow):
self._switch_ui_to_clip()
def next_clip(self):
mask_path = self.out_dir / self.filename.stem / self.cfg.filenames.mask
if mask_path.exists():
mask_path = self._out_path(self.cfg.filenames.mask)
if self._fs_exists(mask_path):
msg = QMessageBox(self)
msg.setWindowTitle("Existing annotation found")
msg.setText(
f"'{self.filename.stem}' already has a saved annotation.\n"
f"'{fsstem(self.filename)}' already has a saved annotation.\n"
"Replace it with your current work, or keep the existing save?"
)
btn_replace = msg.addButton(
@@ -408,13 +473,15 @@ class Annotator(QMainWindow):
)
return
prev_clip = self.selector.clips[idx - 1]
mask_path = self.out_dir / prev_clip.stem / self.cfg.filenames.mask
if not mask_path.exists():
mask_path = fsjoin(self.out_dir, fsstem(prev_clip), self.cfg.filenames.mask)
if not self._fs_exists(mask_path):
QMessageBox.information(
self, "No mask found", f"No saved mask found for '{prev_clip.stem}'."
self,
"No mask found",
f"No saved mask found for '{fsstem(prev_clip)}'.",
)
return
mask_full = np.array(Image.open(mask_path).convert("L"))
mask_full = np.array(self._pil_open(mask_path).convert("L"))
mask = cv2.resize(
(mask_full > 127).astype(np.uint8),
(self.dw, self.dh),