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:
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user