From b773d567233eaa17ab88c4cff79a35704b9f7c53 Mon Sep 17 00:00:00 2001 From: Nathaniel Chen Date: Sat, 4 Jul 2026 18:40:59 -0400 Subject: [PATCH 1/9] feat: add gradio-free core modules (hub, transforms, inference, examples) Carve reusable, gradio-free logic out of the app so a later CLI task and the app can share it: compute_stft (transforms.py), model_infer/warmup/ signal_to_spectrogram (inference.py), HF Hub + local model resolution (hub.py), and a deterministic synthetic demo signal (examples.py). src/tokeye/app/ is untouched; a later task rewires it to delegate here. Adds huggingface-hub and tqdm as core dependencies. --- pyproject.toml | 2 + src/tokeye/examples.py | 84 +++++++++++++++++++++++ src/tokeye/hub.py | 140 +++++++++++++++++++++++++++++++++++++++ src/tokeye/inference.py | 54 +++++++++++++++ src/tokeye/transforms.py | 44 ++++++++++++ tests/test_examples.py | 48 ++++++++++++++ tests/test_hub.py | 67 +++++++++++++++++++ tests/test_inference.py | 51 ++++++++++++++ tests/test_transforms.py | 57 ++++++++++++++++ uv.lock | 6 +- 10 files changed, 552 insertions(+), 1 deletion(-) create mode 100644 src/tokeye/examples.py create mode 100644 src/tokeye/hub.py create mode 100644 src/tokeye/inference.py create mode 100644 src/tokeye/transforms.py create mode 100644 tests/test_examples.py create mode 100644 tests/test_hub.py create mode 100644 tests/test_inference.py create mode 100644 tests/test_transforms.py diff --git a/pyproject.toml b/pyproject.toml index c12aacf..87f5c5e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,8 @@ dependencies = [ "omegaconf", "tables", "pydantic", + "huggingface-hub>=0.30", + "tqdm", ] [project.optional-dependencies] diff --git a/src/tokeye/examples.py b/src/tokeye/examples.py new file mode 100644 index 0000000..d1c16c4 --- /dev/null +++ b/src/tokeye/examples.py @@ -0,0 +1,84 @@ +"""Deterministic synthetic example signal for demos and smoke tests. + +Produces a 1D signal with Gaussian noise plus a handful of sustained +frequency-ramp "chirps" (coherent activity) and short Gaussian-windowed +"bursts" (transient activity), so a log1p spectrogram of the output shows +both activity classes TokEye is trained to segment. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from scipy.signal import chirp +from scipy.signal.windows import tukey + +NOISE_SIGMA = 1.0 +CHIRP_AMPLITUDE = 3.0 +BURST_AMPLITUDE = 5.0 + + +def _add_chirp( + t: np.ndarray, fs: float, duration_s: float, rng: np.random.Generator +) -> np.ndarray: + """A sustained linear/quadratic frequency ramp over part of the signal.""" + seg_duration = rng.uniform(0.3, 0.6) * duration_s + start_time = rng.uniform(0.0, duration_s - seg_duration) + mask = (t >= start_time) & (t < start_time + seg_duration) + + local_t = t[mask] - start_time + f0 = rng.uniform(0.02, 0.10) * fs + f1 = rng.uniform(0.10, 0.30) * fs + method = rng.choice(["linear", "quadratic"]) + ramp = chirp(local_t, f0=f0, f1=f1, t1=seg_duration, method=method) + envelope = tukey(local_t.size, alpha=0.1) if local_t.size else local_t + + out = np.zeros_like(t) + out[mask] = CHIRP_AMPLITUDE * ramp * envelope + return out + + +def _add_burst( + t: np.ndarray, fs: float, duration_s: float, rng: np.random.Generator +) -> np.ndarray: + """A short Gaussian-windowed tone burst (transient activity).""" + burst_duration = rng.uniform(0.002, 0.01) * duration_s + center_time = rng.uniform(0.0, duration_s) + carrier_freq = rng.uniform(0.05, 0.40) * fs + phase = rng.uniform(0.0, 2 * np.pi) + + sigma = burst_duration / 6.0 + envelope = np.exp(-0.5 * ((t - center_time) / sigma) ** 2) + return BURST_AMPLITUDE * envelope * np.sin(2 * np.pi * carrier_freq * t + phase) + + +def make_example_signal( + duration_s: float = 2.0, + fs: float = 200_000.0, + seed: int = 0, +) -> np.ndarray: + rng = np.random.default_rng(seed) + n_samples = int(duration_s * fs) + t = np.arange(n_samples, dtype=np.float64) / fs + + sig = rng.normal(0.0, NOISE_SIGMA, size=n_samples) + + for _ in range(rng.integers(2, 4)): # 2-3 chirps + sig += _add_chirp(t, fs, duration_s, rng) + + for _ in range(rng.integers(2, 4)): # 2-3 bursts + sig += _add_burst(t, fs, duration_s, rng) + + return sig.astype(np.float32) + + +def write_example_signal( + directory: Path, + filename: str = "tokeye_example.npy", +) -> Path: + directory = Path(directory) + directory.mkdir(parents=True, exist_ok=True) + out_path = directory / filename + np.save(out_path, make_example_signal()) + return out_path diff --git a/src/tokeye/hub.py b/src/tokeye/hub.py new file mode 100644 index 0000000..4c03ea3 --- /dev/null +++ b/src/tokeye/hub.py @@ -0,0 +1,140 @@ +"""Model loading: HuggingFace Hub downloads plus local file resolution. + +Gradio-free so it can be shared between the app and a future CLI. +""" + +from __future__ import annotations + +import logging +import os +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from huggingface_hub import hf_hub_download + +from .models.big_tf_unet.config_big_tf_unet import BigTFUNetConfig +from .models.big_tf_unet.model_big_tf_unet import BigTFUNetModel + +if TYPE_CHECKING: + from collections.abc import Callable + +logger = logging.getLogger(__name__) + +DEFAULT_REPO_ID = os.environ.get("TOKEYE_HF_REPO", "PlasmaControl/tokeye") +DEFAULT_MODEL = "big_tf_unet" + +_PATH_SUFFIXES = {".pt", ".pt2"} + + +@dataclass(frozen=True) +class ModelSpec: + name: str + filename: str # file in the HF repo + builder: Callable[[], nn.Module] + + +MODEL_REGISTRY: dict[str, ModelSpec] = { + "big_tf_unet": ModelSpec( + "big_tf_unet", + "big_tf_unet_251210.pt", + lambda: BigTFUNetModel(BigTFUNetConfig()), + ), +} + + +def resolve_device(device: str = "auto") -> str: + if device == "auto": + return "cuda" if torch.cuda.is_available() else "cpu" + return device + + +def download_model(name: str = DEFAULT_MODEL, repo_id: str | None = None) -> Path: + try: + spec = MODEL_REGISTRY[name] + except KeyError as exc: + raise ValueError( + f"Unknown model {name!r}; valid names: {sorted(MODEL_REGISTRY)}" + ) from exc + resolved_repo_id = repo_id or DEFAULT_REPO_ID + return Path(hf_hub_download(resolved_repo_id, spec.filename)) + + +def _build_from_state_dict(state_dict: Mapping, device: str) -> nn.Module: + mismatches: list[str] = [] + for spec in MODEL_REGISTRY.values(): + model = spec.builder() + try: + model.load_state_dict(state_dict, strict=True) + except RuntimeError as exc: + mismatches.append(f"{spec.name}: {exc}") + continue + return model.to(device).eval() + + details = "\n".join(mismatches) + raise ValueError( + "State dict does not match any known TokEye architecture " + f"({', '.join(sorted(MODEL_REGISTRY))}).\n{details}" + ) + + +def _load_from_registry(name: str, device: str) -> nn.Module: + spec = MODEL_REGISTRY[name] + path = download_model(name) + state_dict = torch.load(path, map_location=device, weights_only=True) + model = spec.builder() + model.load_state_dict(state_dict) + return model.to(device).eval() + + +def _load_pt2(path: Path, device: str) -> nn.Module: + module = torch.export.load(str(path)).module() + return module.to(device) + + +def _load_pt(path: Path, device: str) -> nn.Module: + try: + loaded = torch.load(path, map_location=device, weights_only=True) + except Exception: + # Legacy checkpoint pickled as a full module (not just a state dict). + # Only ever done for local files: the registry/download path above + # always loads with weights_only=True. + logger.warning( + "%s could not be loaded safely (weights_only=True); falling back " + "to unpickling the full file. Only do this for local files you " + "trust.", + path, + ) + model = torch.load(path, map_location=device, weights_only=False) + return model.to(device).eval() + + if isinstance(loaded, Mapping): + return _build_from_state_dict(loaded, device) + + return loaded.to(device).eval() + + +def load_model(source: str | Path = DEFAULT_MODEL, device: str = "auto") -> nn.Module: + resolved_device = resolve_device(device) + name = str(source) + + if name in MODEL_REGISTRY: + return _load_from_registry(name, resolved_device) + + path = Path(source) + if not path.exists(): + if path.suffix in _PATH_SUFFIXES: + raise FileNotFoundError(f"Model file not found: {path}") + raise ValueError( + f"Unknown model {name!r}; valid registry names: {sorted(MODEL_REGISTRY)}" + ) + + if path.suffix == ".pt2": + return _load_pt2(path, resolved_device) + if path.suffix == ".pt": + return _load_pt(path, resolved_device) + + raise ValueError(f"Unsupported model file suffix: {path.suffix!r}") diff --git a/src/tokeye/inference.py b/src/tokeye/inference.py new file mode 100644 index 0000000..f074e6d --- /dev/null +++ b/src/tokeye/inference.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import logging + +import numpy as np +import torch +import torch.nn as nn +from tqdm.auto import tqdm + +from .transforms import compute_stft + +logger = logging.getLogger(__name__) + +WARMUP_INPUT_SHAPE = (1, 1, 512, 512) # (batch_size, channels, height, width) + + +def model_infer( + inp_array: np.ndarray | None, + model: nn.Module | torch.export.ExportedProgram | None, +) -> np.ndarray | None: + if inp_array is None or model is None: + logger.warning("Missing input or model for inference") + return None + + logger.info(f"Running inference on input shape: {inp_array.shape}") + + device = next(model.parameters()).device + inp_array = (inp_array - inp_array.mean()) / (inp_array.std() + 1e-6) + inp_tensor = torch.from_numpy(inp_array) + inp_tensor = inp_tensor.unsqueeze(0).unsqueeze(0).float() + inp_tensor = inp_tensor.to(device) + + with torch.no_grad(): + out_tensor = model(inp_tensor) + out_tensor = out_tensor[0] + + out_tensor = torch.sigmoid(out_tensor) + out_tensor = out_tensor.squeeze(0).squeeze(0).cpu() + return out_tensor.numpy() + + +def signal_to_spectrogram(signal: np.ndarray, **stft_kwargs) -> np.ndarray: + """Expand a 1D signal to (1, N) and run it through ``compute_stft``.""" + signal_data = np.expand_dims(signal, axis=0) + return compute_stft(signal_data, **stft_kwargs) + + +def warmup(model: nn.Module, iterations: int = 10) -> None: + """Run dummy forward passes to trigger lazy init / kernel autotuning.""" + device = next(model.parameters()).device + dummy_input = torch.randn(*WARMUP_INPUT_SHAPE, device=device, dtype=torch.float32) + with torch.no_grad(): + for _ in tqdm(range(iterations)): + _ = model(dummy_input) diff --git a/src/tokeye/transforms.py b/src/tokeye/transforms.py new file mode 100644 index 0000000..de0ef2b --- /dev/null +++ b/src/tokeye/transforms.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import numpy as np +from scipy import signal + +DEFAULT_N_FFT = 1024 +DEFAULT_HOP = 256 # app UI default; training recipe used 128 +DEFAULT_CLIP_DC = True +DEFAULT_CLIP_LOW = 1.0 +DEFAULT_CLIP_HIGH = 99.0 + + +def compute_stft( + arr: np.ndarray, + n_fft: int = DEFAULT_N_FFT, + hop: int = DEFAULT_HOP, + window: str = "hann", + clip_dc: bool = DEFAULT_CLIP_DC, + fs: float = 1.0, + clip_low: float = DEFAULT_CLIP_LOW, + clip_high: float = DEFAULT_CLIP_HIGH, +) -> np.ndarray: + win = signal.get_window(window, n_fft) + transform = signal.ShortTimeFFT(win=win, hop=hop, fs=fs) + sxx = transform.stft(arr) + + if sxx.shape[0] == 2: + sxx = sxx[0] * np.conj(sxx[1]) + elif sxx.shape[0] == 1: + sxx = sxx[0] + + sxx = np.abs(sxx) + sxx = np.log1p(sxx) + + # DC clipping + if clip_dc: + sxx = sxx[1:, :] + + # Percentile clipping + vmin, vmax = np.percentile( + sxx, + [clip_low, clip_high], + ) + return np.clip(sxx, vmin, vmax) diff --git a/tests/test_examples.py b/tests/test_examples.py new file mode 100644 index 0000000..f965225 --- /dev/null +++ b/tests/test_examples.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +import numpy as np + +from tokeye.examples import make_example_signal, write_example_signal + + +def test_length_matches_duration_and_fs(): + duration_s, fs = 1.5, 1000.0 + sig = make_example_signal(duration_s=duration_s, fs=fs) + assert sig.shape[0] == int(duration_s * fs) + + +def test_dtype_is_float(): + sig = make_example_signal(duration_s=0.5, fs=1000.0) + assert np.issubdtype(sig.dtype, np.floating) + + +def test_all_finite(): + sig = make_example_signal(duration_s=0.5, fs=1000.0) + assert np.all(np.isfinite(sig)) + + +def test_same_seed_identical(): + sig_a = make_example_signal(duration_s=0.5, fs=1000.0, seed=42) + sig_b = make_example_signal(duration_s=0.5, fs=1000.0, seed=42) + np.testing.assert_array_equal(sig_a, sig_b) + + +def test_different_seed_differs(): + sig_a = make_example_signal(duration_s=0.5, fs=1000.0, seed=0) + sig_b = make_example_signal(duration_s=0.5, fs=1000.0, seed=1) + assert not np.array_equal(sig_a, sig_b) + + +def test_write_example_signal_creates_file(tmp_path): + target_dir = tmp_path / "nested" / "output" + out_path = write_example_signal(target_dir) + assert out_path.exists() + assert out_path.parent == target_dir + loaded = np.load(out_path) + np.testing.assert_array_equal(loaded, make_example_signal()) + + +def test_write_example_signal_custom_filename(tmp_path): + out_path = write_example_signal(tmp_path, filename="custom.npy") + assert out_path.name == "custom.npy" + assert out_path.exists() diff --git a/tests/test_hub.py b/tests/test_hub.py new file mode 100644 index 0000000..0615d8c --- /dev/null +++ b/tests/test_hub.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import logging + +import pytest +import torch +import torch.nn as nn + +from tokeye.hub import DEFAULT_MODEL, MODEL_REGISTRY, load_model + + +def test_default_model_is_registered(): + assert DEFAULT_MODEL in MODEL_REGISTRY + + +def test_load_model_from_registry_downloads_and_loads(tmp_path, monkeypatch): + spec = MODEL_REGISTRY["big_tf_unet"] + weights_path = tmp_path / spec.filename + torch.save(spec.builder().state_dict(), weights_path) + + def fake_hf_hub_download(repo_id, filename, **kwargs): + assert filename == spec.filename + return str(weights_path) + + monkeypatch.setattr("tokeye.hub.hf_hub_download", fake_hf_hub_download) + + model = load_model("big_tf_unet", device="cpu") + + assert not model.training # eval mode + with torch.no_grad(): + out = model(torch.randn(1, 1, 64, 64)) + assert out[0].shape == (1, 2, 64, 64) + + +def test_load_model_from_local_path_with_matching_state_dict(tmp_path): + spec = MODEL_REGISTRY["big_tf_unet"] + weights_path = tmp_path / "checkpoint.pt" + torch.save(spec.builder().state_dict(), weights_path) + + model = load_model(weights_path, device="cpu") + + assert not model.training + with torch.no_grad(): + out = model(torch.randn(1, 1, 64, 64)) + assert out[0].shape == (1, 2, 64, 64) + + +def test_load_model_legacy_pickled_module_falls_back(tmp_path, caplog): + legacy_path = tmp_path / "legacy.pt" + torch.save(nn.Linear(2, 2), legacy_path) + + with caplog.at_level(logging.WARNING, logger="tokeye.hub"): + model = load_model(legacy_path, device="cpu") + + assert isinstance(model, nn.Linear) + assert not model.training + assert "trust" in caplog.text.lower() + + +def test_load_model_unknown_name_raises_value_error_listing_registry(): + with pytest.raises(ValueError, match="big_tf_unet"): + load_model("not_a_real_model_name") + + +def test_load_model_nonexistent_path_raises_file_not_found(): + with pytest.raises(FileNotFoundError): + load_model("/nonexistent/directory/does_not_exist.pt") diff --git a/tests/test_inference.py b/tests/test_inference.py new file mode 100644 index 0000000..8ca6e5c --- /dev/null +++ b/tests/test_inference.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +import numpy as np +import torch +import torch.nn as nn + +from tokeye.inference import model_infer, signal_to_spectrogram, warmup + + +class _IdentityLikeModel(nn.Module): + """Returns a 1-tuple of (B, 2, H, W), mimicking BigTFUNetModel's output.""" + + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv2d(1, 2, kernel_size=1) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor]: + return (self.conv(x),) + + +def test_model_infer_output_shape_and_range(): + model = _IdentityLikeModel() + model.eval() + inp = np.random.default_rng(0).normal(size=(64, 64)).astype(np.float32) + + out = model_infer(inp, model) + + assert out.shape == (2, 64, 64) + assert np.all(out >= 0.0) and np.all(out <= 1.0) # sigmoid range + + +def test_model_infer_none_input_returns_none(): + model = _IdentityLikeModel() + assert model_infer(None, model) is None + + +def test_model_infer_none_model_returns_none(): + inp = np.zeros((32, 32), dtype=np.float32) + assert model_infer(inp, None) is None + + +def test_signal_to_spectrogram_returns_2d(): + signal = np.random.default_rng(0).normal(size=4096).astype(np.float64) + spec = signal_to_spectrogram(signal, n_fft=256, hop=64) + assert spec.ndim == 2 + + +def test_warmup_runs_without_error(): + model = _IdentityLikeModel() + model.eval() + warmup(model, iterations=2) # should not raise diff --git a/tests/test_transforms.py b/tests/test_transforms.py new file mode 100644 index 0000000..06768ab --- /dev/null +++ b/tests/test_transforms.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import numpy as np + +from tokeye.transforms import compute_stft + + +def _sine_signal(n_samples: int = 4096, freq: float = 0.05) -> np.ndarray: + t = np.arange(n_samples) + return np.sin(2 * np.pi * freq * t).astype(np.float64)[np.newaxis, :] + + +def test_compute_stft_output_is_2d(): + arr = _sine_signal() + out = compute_stft(arr, n_fft=256, hop=64) + assert out.ndim == 2 + + +def test_clip_dc_removes_one_frequency_row(): + arr = _sine_signal() + with_dc_clip = compute_stft(arr, n_fft=256, hop=64, clip_dc=True) + without_dc_clip = compute_stft(arr, n_fft=256, hop=64, clip_dc=False) + assert with_dc_clip.shape[0] == without_dc_clip.shape[0] - 1 + assert with_dc_clip.shape[1] == without_dc_clip.shape[1] + + +def test_expected_frame_count(): + # For n_samples=4096, n_fft=1024, hop=256, scipy's ShortTimeFFT (with its + # edge-padding convention) yields exactly 19 frames; pin that number so a + # change to the padding/framing convention is caught as a regression. + n_samples = 4096 + n_fft = 1024 + hop = 256 + arr = _sine_signal(n_samples=n_samples) + out = compute_stft(arr, n_fft=n_fft, hop=hop) + assert out.shape[1] == 19 + + +def test_values_within_percentile_clip_bounds(): + # A signal with an extreme outlier burst has a much wider raw dynamic + # range than its 1st-99th percentile band; percentile clipping should + # compress the output range to (approximately) that narrower band. + rng = np.random.default_rng(0) + arr = rng.normal(size=4096) + arr[2000] = 1000.0 # single extreme outlier + arr = arr[np.newaxis, :].astype(np.float64) + + clip_low, clip_high = 1.0, 99.0 + clipped = compute_stft( + arr, n_fft=256, hop=64, clip_low=clip_low, clip_high=clip_high + ) + unclipped = compute_stft(arr, n_fft=256, hop=64, clip_low=0.0, clip_high=100.0) + + assert np.all(np.isfinite(clipped)) + assert clipped.min() >= unclipped.min() + assert clipped.max() <= unclipped.max() + assert clipped.max() < unclipped.max() # outlier should be visibly clipped diff --git a/uv.lock b/uv.lock index 0cfd854..4b6a0bd 100644 --- a/uv.lock +++ b/uv.lock @@ -2707,10 +2707,11 @@ wheels = [ [[package]] name = "tokeye" -version = "0.9.0" +version = "0.9.5" source = { editable = "." } dependencies = [ { name = "gradio" }, + { name = "huggingface-hub" }, { name = "matplotlib" }, { name = "numpy" }, { name = "omegaconf" }, @@ -2721,6 +2722,7 @@ dependencies = [ { name = "torchcodec" }, { name = "torchinfo" }, { name = "torchvision" }, + { name = "tqdm" }, ] [package.optional-dependencies] @@ -2769,6 +2771,7 @@ train = [ requires-dist = [ { name = "gradio" }, { name = "h5py", marker = "extra == 'train'" }, + { name = "huggingface-hub", specifier = ">=0.30" }, { name = "ipykernel", marker = "extra == 'train'" }, { name = "ipywidgets", marker = "extra == 'train'" }, { name = "joblib", marker = "extra == 'train'" }, @@ -2790,6 +2793,7 @@ requires-dist = [ { name = "torchinfo" }, { name = "torchmetrics", marker = "extra == 'train'" }, { name = "torchvision" }, + { name = "tqdm" }, { name = "wavio", marker = "extra == 'train'" }, ] provides-extras = ["train"] From a1b42f3db2808814dc6ba8788ca7aa4ddaa901f2 Mon Sep 17 00:00:00 2001 From: Nathaniel Chen Date: Sat, 4 Jul 2026 18:59:40 -0400 Subject: [PATCH 2/9] feat: guided onboarding in analyze tab; package logo; delete dead modules - analyze.py: default-populated transform_args (kills the silent-None trap where nothing works until Setup Transform is clicked), registry-aware model dropdown with automatic HF download + cache-aware toast, one-click "Analyze" button (ensure_model -> infer -> visualize), "Load Example Signal" button backed by tokeye.examples, friendlier copy. - load.py: delegates model loading/warmup to tokeye.hub/tokeye.inference, compute_stft to tokeye.transforms; deletes the now-redundant torch.load branching and inline warmup loop. - visualize.py: enhance() copies its input before mutating it in place, so repeated Visualize clicks no longer corrupt the stored inference output. - __main__.py: extracts main(port, share, open_browser), replaces ad-hoc sys.argv parsing with argparse, loads the logo via importlib.resources (guarded, no crash if missing) now that it's packaged under src/tokeye/app/assets/. - Deletes dead modules superseded by the Task 1 core package: app/analyze /transforms.py, app/processing/, app/utils/analyze.py, tokeye/analysis/ (all already had stale TokEye.* import paths from before the package rename). - Adds tests/test_app_load.py covering model_load delegation, find_models, load_single 1D/2D handling, the new is_model_cached helper, and a create_app() smoke-construction test. --- README.md | 2 +- src/tokeye/analysis/__init__.py | 4 - src/tokeye/analysis/batch_analysis.py | 240 ----------- src/tokeye/app/__main__.py | 81 ++-- src/tokeye/app/analyze/analyze.py | 140 ++++++- src/tokeye/app/analyze/load.py | 146 +++---- src/tokeye/app/analyze/transforms.py | 37 -- src/tokeye/app/analyze/visualize.py | 1 + {assets => src/tokeye/app/assets}/logo.png | Bin src/tokeye/app/processing/QUICKSTART.md | 310 -------------- src/tokeye/app/processing/README.md | 459 --------------------- src/tokeye/app/processing/__init__.py | 0 src/tokeye/app/processing/inference.py | 66 --- src/tokeye/app/processing/postprocess.py | 386 ----------------- src/tokeye/app/processing/tiling.py | 292 ------------- src/tokeye/app/utils/analyze.py | 391 ------------------ tests/test_app_load.py | 87 ++++ 17 files changed, 337 insertions(+), 2305 deletions(-) delete mode 100644 src/tokeye/analysis/__init__.py delete mode 100644 src/tokeye/analysis/batch_analysis.py delete mode 100644 src/tokeye/app/analyze/transforms.py rename {assets => src/tokeye/app/assets}/logo.png (100%) delete mode 100644 src/tokeye/app/processing/QUICKSTART.md delete mode 100644 src/tokeye/app/processing/README.md delete mode 100644 src/tokeye/app/processing/__init__.py delete mode 100644 src/tokeye/app/processing/inference.py delete mode 100644 src/tokeye/app/processing/postprocess.py delete mode 100644 src/tokeye/app/processing/tiling.py delete mode 100644 src/tokeye/app/utils/analyze.py create mode 100644 tests/test_app_load.py diff --git a/README.md b/README.md index 7aa873e..7490fa5 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,5 @@

- TokEye Logo + TokEye Logo

# TokEye diff --git a/src/tokeye/analysis/__init__.py b/src/tokeye/analysis/__init__.py deleted file mode 100644 index 1a0ee1b..0000000 --- a/src/tokeye/analysis/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .batch_analysis import default_settings, run_batch_analysis - -__all__ = ["run_batch_analysis", "default_settings"] - diff --git a/src/tokeye/analysis/batch_analysis.py b/src/tokeye/analysis/batch_analysis.py deleted file mode 100644 index 1d82247..0000000 --- a/src/tokeye/analysis/batch_analysis.py +++ /dev/null @@ -1,240 +0,0 @@ -import logging -import time -from pathlib import Path - -import numpy as np -import torch -import torch.nn as nn - -# Import model loading function from the analyze module -from TokEye.app.analyze.load import model_load -from tqdm.auto import tqdm - -# Configure logging -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", -) -logger = logging.getLogger(__name__) - -# Default settings -default_settings = { - "model_path": "data/models/big_mode_v1.pth", - "input_path": "data/batch_inputs", - "output_path": "data/batch_outputs", - "analysis_mode": "amplitude", - "threshold": 0.5, - "polling_interval": 5, # seconds between directory scans -} - - -def load_spectrogram(filepath: Path) -> np.ndarray | None: - """Load a 2D spectrogram from a numpy file. - - Args: - filepath: Path to the .npy file containing a 2D spectrogram - - Returns: - 2D numpy array or None if loading fails - """ - try: - spec = np.load(filepath) - if spec.ndim != 2: - logger.error(f"Expected 2D array, got {spec.ndim}D: {filepath.name}") - return None - if spec.size == 0: - logger.error(f"Empty array: {filepath.name}") - return None - return spec - except Exception as e: - logger.error(f"Failed to load {filepath.name}: {e}") - return None - - -def run_inference( - spec: np.ndarray, - model: nn.Module | torch.export.ExportedProgram, -) -> np.ndarray | None: - """Run inference on a 2D spectrogram. - - Args: - spec: 2D numpy array (H, W) - model: Loaded PyTorch model - - Returns: - 3D numpy array (2, H, W) with two output channels or None if inference fails - """ - try: - device = next(model.parameters()).device - - # Normalize input - spec_norm = (spec - spec.mean()) / (spec.std() + 1e-6) - - # Convert to tensor: (1, 1, H, W) - inp_tensor = torch.from_numpy(spec_norm) - inp_tensor = inp_tensor.unsqueeze(0).unsqueeze(0).float() - inp_tensor = inp_tensor.to(device) - - # Run inference - with torch.no_grad(): - out_tensor = model(inp_tensor) - - # Remove batch dimension: (1, 2, H, W) -> (2, H, W) - out_tensor = out_tensor[0] - - # Apply sigmoid activation - out_tensor = torch.sigmoid(out_tensor) - - # Convert to numpy - return out_tensor.cpu().numpy() - - except Exception as e: - logger.error(f"Inference failed: {e}") - return None - - -def process_file( - input_path: Path, - output_path: Path, - model: nn.Module | torch.export.ExportedProgram, -) -> bool: - """Process a single file: load, infer, save. - - Args: - input_path: Path to input .npy file - output_path: Path to output .npy file - model: Loaded PyTorch model - - Returns: - True if successful, False otherwise - """ - # Load spectrogram - spec = load_spectrogram(input_path) - if spec is None: - return False - - # Run inference - result = run_inference(spec, model) - if result is None: - return False - - # Save result - try: - output_path.parent.mkdir(parents=True, exist_ok=True) - np.save(output_path, result) - return True - except Exception as e: - logger.error(f"Failed to save {output_path.name}: {e}") - return False - - -def scan_directory(input_dir: Path) -> set[str]: - """Scan directory for .npy files. - - Args: - input_dir: Directory to scan - - Returns: - Set of filenames (not full paths) - """ - if not input_dir.exists(): - return set() - - npy_files = input_dir.glob("*.npy") - return {f.name for f in npy_files} - - -def run_batch_analysis(settings: dict): - """Main batch analysis loop. - - Continuously monitors input directory for new spectrogram files, - runs inference, and saves results. Runs until interrupted with Ctrl+C. - - Args: - settings: Dictionary with configuration: - - model_path: Path to model file - - input_path: Input directory to monitor - - output_path: Output directory for results - - polling_interval: Seconds between directory scans - """ - # Extract settings - model_path = Path(settings["model_path"]) - input_dir = Path(settings["input_path"]) - output_dir = Path(settings["output_path"]) - polling_interval = settings.get("polling_interval", 5) - - # Ensure directories exist - input_dir.mkdir(parents=True, exist_ok=True) - output_dir.mkdir(parents=True, exist_ok=True) - - # Load model - logger.info("=" * 60) - logger.info("Starting Batch Analysis System") - logger.info("=" * 60) - logger.info(f"Model: {model_path}") - logger.info(f"Input directory: {input_dir}") - logger.info(f"Output directory: {output_dir}") - logger.info(f"Polling interval: {polling_interval}s") - logger.info("=" * 60) - - try: - model = model_load(model_path) - except Exception as e: - logger.error(f"Failed to load model: {e}") - return - - # Initialize tracking - processed_files = set() - - logger.info("Monitoring for new files... (Press Ctrl+C to stop)") - logger.info("") - - # Main loop - try: - while True: - # Scan for all files in directory - current_files = scan_directory(input_dir) - - # Update tracking: remove deleted files - deleted_files = processed_files - current_files - if deleted_files: - for fname in deleted_files: - processed_files.discard(fname) - logger.info(f"File deleted, will reprocess if added again: {fname}") - - # Find new files to process - new_files = current_files - processed_files - - if new_files: - logger.info(f"Found {len(new_files)} new file(s)") - - # Process files with progress bar - new_files_list = sorted(new_files) - for fname in tqdm(new_files_list, desc="Processing", unit="file"): - input_path = input_dir / fname - output_path = output_dir / fname - - # Process the file - success = process_file(input_path, output_path, model) - - if success: - processed_files.add(fname) - logger.info(f"✓ Processed: {fname}") - else: - logger.warning(f"✗ Failed: {fname}") - - logger.info("") - - # Sleep before next scan - time.sleep(polling_interval) - - except KeyboardInterrupt: - logger.info("") - logger.info("=" * 60) - logger.info("Shutting down gracefully...") - logger.info(f"Total files processed: {len(processed_files)}") - logger.info("=" * 60) - - -if __name__ == "__main__": - run_batch_analysis(default_settings) diff --git a/src/tokeye/app/__main__.py b/src/tokeye/app/__main__.py index 54720f2..12811a5 100644 --- a/src/tokeye/app/__main__.py +++ b/src/tokeye/app/__main__.py @@ -2,8 +2,11 @@ TokEye Main Inference """ +from __future__ import annotations + +import argparse +import importlib.resources import logging -import sys from pathlib import Path import gradio as gr @@ -23,8 +26,7 @@ logging.getLogger("uvicorn").setLevel(logging.WARNING) logging.getLogger("httpx").setLevel(logging.WARNING) -# Current working directory -cwd = Path.cwd() +logger = logging.getLogger(__name__) def create_app() -> gr.Blocks: @@ -33,14 +35,16 @@ def create_app() -> gr.Blocks: theme=make_theme(), css="footer{display:none !important}", ) as app: - gr.Image( - str(Path.cwd() / "assets" / "logo.png"), - height=300, - interactive=False, - container=False, - show_download_button=False, - show_fullscreen_button=False, - ) + logo_path = importlib.resources.files("tokeye.app").joinpath("assets/logo.png") + if logo_path.is_file(): + gr.Image( + str(logo_path), + height=300, + interactive=False, + container=False, + show_download_button=False, + show_fullscreen_button=False, + ) with gr.Tab("Analyze"): analyze_tab() with gr.Tab("Annotate"): @@ -50,35 +54,44 @@ def create_app() -> gr.Blocks: return app -def get_port(): - if "--port" in sys.argv: - port_index = sys.argv.index("--port") + 1 - if port_index < len(sys.argv): - return int(sys.argv[port_index]) - return DEFAULT_PORT - - -def launch(app, port): - app.launch( - # favicon_path="assets/ICON.ico", # Set up favicon later - share="--share" in sys.argv, - inbrowser="--open" in sys.argv, - server_port=port, - ) - - -if __name__ == "__main__": - logging.info(f"Initializing TokEye in: {cwd}") - # Set up +def main( + port: int = DEFAULT_PORT, + share: bool = False, + open_browser: bool = False, +) -> None: + logger.info(f"Initializing TokEye in: {Path.cwd()}") app = create_app() - # Launch application - port = get_port() for _ in range(MAX_PORT_ATTEMPTS): try: - launch(app, port) + app.launch( + share=share, + inbrowser=open_browser, + server_port=port, + ) except OSError: print(f"Failed on port {port}") port -= 1 except Exception as error: print(f"{error}") break + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + prog="python -m tokeye.app", + description="Launch the TokEye Gradio app.", + ) + parser.add_argument( + "--port", type=int, default=DEFAULT_PORT, help="Port to serve the app on." + ) + parser.add_argument( + "--share", action="store_true", help="Create a public Gradio share link." + ) + parser.add_argument( + "--open", + dest="open_browser", + action="store_true", + help="Open the app in a browser on launch.", + ) + args = parser.parse_args() + main(port=args.port, share=args.share, open_browser=args.open_browser) diff --git a/src/tokeye/app/analyze/analyze.py b/src/tokeye/app/analyze/analyze.py index 67ead3c..30a01bd 100644 --- a/src/tokeye/app/analyze/analyze.py +++ b/src/tokeye/app/analyze/analyze.py @@ -1,11 +1,24 @@ +from __future__ import annotations + import logging from pathlib import Path import gradio as gr +from tokeye.hub import DEFAULT_MODEL, MODEL_REGISTRY +from tokeye.transforms import ( + DEFAULT_CLIP_DC, + DEFAULT_CLIP_HIGH, + DEFAULT_CLIP_LOW, + DEFAULT_HOP, + DEFAULT_N_FFT, +) + from .load import ( find_models, find_signals, + is_model_cached, + load_example_signal, load_multi, load_single, model_infer, @@ -31,7 +44,7 @@ def setup_stft_transform(n_fft, hop_length, clip_dc, clip_low, clip_high): def refresh_dropdowns(signal_directory): - models = find_models() + models = list(MODEL_REGISTRY) + find_models() signals = find_signals(signal_directory) return [ gr.Dropdown(choices=models), @@ -59,10 +72,35 @@ def toggle_view_groups(mode): def wrapper_model_load(model_file): - """Wrapper to convert string path to Path object for model_load.""" + """Load a model by registry name or local path. + + Registry names are downloaded from Hugging Face on first use (and cached + thereafter); local paths are loaded directly. Failures are reported via + a toast instead of raising, so a bad model choice degrades to "nothing + loaded" rather than crashing the pipeline. + """ if not model_file: return None - return model_load(Path(model_file)) + try: + if model_file in MODEL_REGISTRY and not is_model_cached(model_file): + gr.Info(f"Downloading {model_file} from Hugging Face (~30 MB, one-time)…") + return model_load(model_file) + except Exception as e: + gr.Warning(f"Model load failed: {e}") + return None + + +def ensure_model(model, model_file): + """Load the model on first use; pass an already-loaded model through.""" + if model is not None: + return model + return wrapper_model_load(model_file) + + +def warn_if_missing_signal(signal_transform): + """Nudge the user if Analyze is clicked before any signal is loaded.""" + if signal_transform is None: + gr.Warning("Load a signal first (or click Load Example Signal)") def wrapper_load_single(signal_directory, signal_file, transform_args): @@ -82,6 +120,13 @@ def wrapper_load_multi(signal_directory, signal_1, signal_2, transform_args): ) +def wrapper_load_example(transform_args): + """Wrapper to guard against missing transform state.""" + if transform_args is None: + return None + return load_example_signal(transform_args) + + def analyze_tab(): # User Interface with gr.Column(): @@ -92,8 +137,13 @@ def analyze_tab(): with gr.Group(): model_file = gr.Dropdown( label="Analysis Model", - info="Select Model For Analysis", - choices=find_models(), + info=( + "Built-in models download automatically from Hugging Face " + "on first load (~30 MB, cached). Local model/*.pt files " + "also listed." + ), + choices=list(MODEL_REGISTRY) + find_models(), + value=DEFAULT_MODEL, interactive=True, allow_custom_value=True, ) @@ -101,21 +151,29 @@ def analyze_tab(): ## Transform with gr.Group(): with gr.Group(): - clip_low_sld = gr.Slider(0, 100, value=1, step=1, label="% Clip Low") - clip_high_sld = gr.Slider(0, 100, value=99, step=1, label="% Clip High") + clip_low_sld = gr.Slider( + 0, 100, value=DEFAULT_CLIP_LOW, step=1, label="% Clip Low" + ) + clip_high_sld = gr.Slider( + 0, 100, value=DEFAULT_CLIP_HIGH, step=1, label="% Clip High" + ) with gr.Tab("STFT"): n_fft = gr.Slider( - 256, 2048, value=1024, step=256, label="Number of Bins" + 256, 2048, value=DEFAULT_N_FFT, step=256, label="Number of Bins" ) - hop_length = gr.Slider(64, 512, value=256, step=64, label="Hop Size") - clip_dc = gr.Checkbox(value=True, label="Remove DC (Bottom) Bin") - setup_tranform_stft_btn = gr.Button("Setup Transform") + hop_length = gr.Slider( + 64, 512, value=DEFAULT_HOP, step=64, label="Hop Size" + ) + clip_dc = gr.Checkbox( + value=DEFAULT_CLIP_DC, label="Remove DC (Bottom) Bin" + ) + setup_tranform_stft_btn = gr.Button("Apply Transform Settings") ## Signal Directory signal_directory = gr.Textbox( label="Signal Directory", value="data/input", - info="Directory containing shot subdirectories", + info="Directory containing .npy signal files", ) with gr.Tab("Single Signal Input"), gr.Column(): @@ -127,6 +185,7 @@ def analyze_tab(): allow_custom_value=True, ) load_single_btn = gr.Button("Load Signal") + load_example_btn = gr.Button("Load Example Signal") # Multi Signal with gr.Tab("Cross Signal Input"), gr.Column(): @@ -175,6 +234,7 @@ def analyze_tab(): with gr.Group(visible=False) as mask_grp, gr.Column(): threshold_sld = gr.Slider(0, 1, value=0.5, step=0.01, label="Threshold") + analyze_btn = gr.Button("Analyze", variant="primary") visualize_btn = gr.Button("Visualize") visualize_out = gr.Image(label="Visualization", type="pil") @@ -182,7 +242,15 @@ def analyze_tab(): model = gr.State() signal_transform = gr.State() inference_output = gr.State() - transform_args = gr.State() + transform_args = gr.State( + setup_stft_transform( + DEFAULT_N_FFT, + DEFAULT_HOP, + DEFAULT_CLIP_DC, + DEFAULT_CLIP_LOW, + DEFAULT_CLIP_HIGH, + ) + ) # Event Handling ## Refresh Page @@ -265,6 +333,24 @@ def analyze_tab(): ], outputs=[extract_out], ) + load_example_btn.click( + fn=wrapper_load_example, + inputs=[transform_args], + outputs=[signal_transform], + ).then( + fn=show_image, + inputs=[ + gr.State("Original"), + signal_transform, + inference_output, + gr.State(False), + gr.State(False), + vmin_sld, + vmax_sld, + threshold_sld, + ], + outputs=[extract_out], + ) ## Visualization view_mode.change( @@ -294,3 +380,31 @@ def analyze_tab(): ], outputs=[visualize_out], ) + + ## One-click Analyze: ensure a model is loaded, run inference, visualize + analyze_btn.click( + fn=warn_if_missing_signal, + inputs=[signal_transform], + outputs=[], + ).then( + fn=ensure_model, + inputs=[model, model_file], + outputs=[model], + ).then( + fn=model_infer, + inputs=[signal_transform, model], + outputs=[inference_output], + ).then( + fn=show_image, + inputs=[ + view_mode, + signal_transform, + inference_output, + out_1_chk, + out_2_chk, + vmin_sld, + vmax_sld, + threshold_sld, + ], + outputs=[visualize_out], + ) diff --git a/src/tokeye/app/analyze/load.py b/src/tokeye/app/analyze/load.py index f644186..d740fe6 100644 --- a/src/tokeye/app/analyze/load.py +++ b/src/tokeye/app/analyze/load.py @@ -1,17 +1,41 @@ +from __future__ import annotations + import logging from pathlib import Path +from typing import TYPE_CHECKING import numpy as np -import torch -import torch.nn as nn -from tqdm.auto import tqdm - -from .transforms import compute_stft +from huggingface_hub import try_to_load_from_cache + +from tokeye import hub, inference +from tokeye.examples import make_example_signal +from tokeye.inference import model_infer, signal_to_spectrogram +from tokeye.transforms import ( + DEFAULT_CLIP_DC, + DEFAULT_CLIP_HIGH, + DEFAULT_CLIP_LOW, + DEFAULT_HOP, + DEFAULT_N_FFT, + compute_stft, +) + +if TYPE_CHECKING: + import torch.nn as nn + +__all__ = [ + "find_models", + "find_signals", + "is_model_cached", + "load_example_signal", + "load_multi", + "load_single", + "model_infer", + "model_load", + "signal_load", +] logger = logging.getLogger(__name__) -DUMMY_INPUT_SHAPE = (1, 1, 512, 512) # (batch_size, channels, height, width) -WARMUP_ITERATIONS = 10 MODEL_EXTENSIONS = [".pt", ".pt2"] MODEL_DIR = Path("model") SIGNAL_EXTENSIONS = [".npy"] @@ -46,67 +70,25 @@ def find_signals( # Model Functions +def is_model_cached(name: str) -> bool: + """Check whether a registry model's weights are already in the HF cache.""" + spec = hub.MODEL_REGISTRY[name] + cached = try_to_load_from_cache(hub.DEFAULT_REPO_ID, spec.filename) + return isinstance(cached, str) + + def model_load( - filepath: Path, + source: str | Path, device: str = "auto", -) -> nn.Module | torch.export.ExportedProgram: - if not filepath.exists(): - raise FileNotFoundError(f"Model not found: {filepath}") - - if device == "auto": - device = "cuda" if torch.cuda.is_available() else "cpu" - - logger.info(f"Loading model: {filepath.name} on {device}") - - match filepath.suffix: - case ".pt": - model = torch.load( - str(filepath), - map_location=device, - weights_only=False, - ) - model.eval() - case ".pt2": - module = torch.export.load(str(filepath)) - model = module.module() - model.to(device) - case _: - raise ValueError(f"Unsupported model format: {filepath.suffix}") - - logger.info(f"Warming up model ({WARMUP_ITERATIONS} iterations)...") - dummy_input = torch.randn(*DUMMY_INPUT_SHAPE, device=device, dtype=torch.float32) - with torch.no_grad(): - for _ in tqdm(range(WARMUP_ITERATIONS)): - _ = model(dummy_input) +) -> nn.Module: + """Load a model by registry name or local path, then warm it up.""" + logger.info(f"Loading model: {source}") + model = hub.load_model(source, device) + inference.warmup(model) logger.info("Model ready for inference") return model -def model_infer( - inp_array: np.ndarray | None, - model: nn.Module | torch.export.ExportedProgram | None, -) -> np.ndarray | None: - if inp_array is None or model is None: - logger.warning("Missing input or model for inference") - return None - - logger.info(f"Running inference on input shape: {inp_array.shape}") - - device = next(model.parameters()).device - inp_array = (inp_array - inp_array.mean()) / (inp_array.std() + 1e-6) - inp_tensor = torch.from_numpy(inp_array) - inp_tensor = inp_tensor.unsqueeze(0).unsqueeze(0).float() - inp_tensor = inp_tensor.to(device) - - with torch.no_grad(): - out_tensor = model(inp_tensor) - out_tensor = out_tensor[0] - - out_tensor = torch.sigmoid(out_tensor) - out_tensor = out_tensor.squeeze(0).squeeze(0).cpu() - return out_tensor.numpy() - - # Signal Functions def signal_load(filepath) -> np.ndarray | None: try: @@ -137,11 +119,11 @@ def load_single( logger.info(f"Raw signal shape: {signal.shape}") # Apply STFT transform (generalize later) - n_fft = transform_args.get("n_fft", 1024) - hop = transform_args.get("hop_length", 256) - clip_dc = transform_args.get("clip_dc", True) - clip_low = transform_args.get("percentile_low", 1.0) - clip_high = transform_args.get("percentile_high", 99.0) + n_fft = transform_args.get("n_fft", DEFAULT_N_FFT) + hop = transform_args.get("hop_length", DEFAULT_HOP) + clip_dc = transform_args.get("clip_dc", DEFAULT_CLIP_DC) + clip_low = transform_args.get("percentile_low", DEFAULT_CLIP_LOW) + clip_high = transform_args.get("percentile_high", DEFAULT_CLIP_HIGH) return compute_stft( signal_data, @@ -169,11 +151,11 @@ def load_multi( logger.info(f"Raw signal shape: {signal.shape}") # Apply STFT to both signals (generalize later) - n_fft = transform_args.get("n_fft", 1024) - hop = transform_args.get("hop_length", 256) - clip_dc = transform_args.get("clip_dc", True) - clip_low = transform_args.get("percentile_low", 1.0) - clip_high = transform_args.get("percentile_high", 99.0) + n_fft = transform_args.get("n_fft", DEFAULT_N_FFT) + hop = transform_args.get("hop_length", DEFAULT_HOP) + clip_dc = transform_args.get("clip_dc", DEFAULT_CLIP_DC) + clip_low = transform_args.get("percentile_low", DEFAULT_CLIP_LOW) + clip_high = transform_args.get("percentile_high", DEFAULT_CLIP_HIGH) return compute_stft( signal, @@ -183,3 +165,23 @@ def load_multi( clip_low=clip_low, clip_high=clip_high, ) + + +def load_example_signal(transform_args: dict) -> np.ndarray | None: + """Generate the deterministic demo signal and compute its spectrogram.""" + signal = make_example_signal() + + n_fft = transform_args.get("n_fft", DEFAULT_N_FFT) + hop = transform_args.get("hop_length", DEFAULT_HOP) + clip_dc = transform_args.get("clip_dc", DEFAULT_CLIP_DC) + clip_low = transform_args.get("percentile_low", DEFAULT_CLIP_LOW) + clip_high = transform_args.get("percentile_high", DEFAULT_CLIP_HIGH) + + return signal_to_spectrogram( + signal, + n_fft=n_fft, + hop=hop, + clip_dc=clip_dc, + clip_low=clip_low, + clip_high=clip_high, + ) diff --git a/src/tokeye/app/analyze/transforms.py b/src/tokeye/app/analyze/transforms.py deleted file mode 100644 index 5320f2f..0000000 --- a/src/tokeye/app/analyze/transforms.py +++ /dev/null @@ -1,37 +0,0 @@ -import numpy as np -from scipy import signal - - -def compute_stft( - arr: np.ndarray, - n_fft: int = 1024, - hop: int = 128, - window: str = "hann", - clip_dc: bool = True, - fs: float = 1.0, - clip_low: float = 1.0, - clip_high: float = 99.0, -) -> np.ndarray: - win = signal.get_window(window, n_fft) - transform = signal.ShortTimeFFT(win=win, hop=hop, fs=fs) - sxx = transform.stft(arr) - - if sxx.shape[0] == 2: - sxx = sxx[0] * np.conj(sxx[1]) - elif sxx.shape[0] == 1: - sxx = sxx[0] - - sxx = np.abs(sxx) - sxx = np.log1p(sxx) - - # DC clipping - if clip_dc: - sxx = sxx[1:, :] - - # Percentile clipping - vmin, vmax = np.percentile( - sxx, - [clip_low, clip_high], - ) - return np.clip(sxx, vmin, vmax) - diff --git a/src/tokeye/app/analyze/visualize.py b/src/tokeye/app/analyze/visualize.py index a7b5499..0d9cdd6 100644 --- a/src/tokeye/app/analyze/visualize.py +++ b/src/tokeye/app/analyze/visualize.py @@ -46,6 +46,7 @@ def enhance( Returns: RGB array ready for plotting """ + arr = arr.copy() n_channels, h, w = arr.shape # Clip values diff --git a/assets/logo.png b/src/tokeye/app/assets/logo.png similarity index 100% rename from assets/logo.png rename to src/tokeye/app/assets/logo.png diff --git a/src/tokeye/app/processing/QUICKSTART.md b/src/tokeye/app/processing/QUICKSTART.md deleted file mode 100644 index 5650afb..0000000 --- a/src/tokeye/app/processing/QUICKSTART.md +++ /dev/null @@ -1,310 +0,0 @@ -# TokEye Processing - Quick Start Guide - -## Installation - -```bash -# Install dependencies -cd c:/Users/twoga/Documents/GitHub/PlasmaControlGroup/TokEye -uv pip install scipy opencv-python -``` - -## Basic Usage - -### 1. Signal to Spectrogram - -```python -import numpy as np -from TokEye.processing import apply_preemphasis, compute_stft - -# Load your signal -signal = np.load('plasma_signal.npy') # shape: (N,) - -# Apply preemphasis -emphasized = apply_preemphasis(signal, alpha=0.97) - -# Compute STFT spectrogram -spectrogram = compute_stft( - emphasized, - n_fft=1024, - hop_length=128, - window='hann', - clip_dc=True, - fs=100000, # 100 kHz sampling rate -) -# Output shape: (freq_bins, time_frames) -``` - -### 2. Wavelet Decomposition - -```python -from TokEye.processing import compute_wavelet - -# Compute wavelet decomposition -coeffs = compute_wavelet( - signal, - wavelet='db8', - level=9, - mode='sym', - order='freq', -) -# Output shape: (512, coeffs_per_node) where 512 = 2^9 -``` - -### 3. Model Inference Pipeline - -```python -from TokEye.processing import ( - tile_spectrogram, - load_model, - batch_inference, - stitch_predictions, -) - -# Pad spectrogram to tile size (256x256) -import numpy as np -if spectrogram.shape[0] != 256: - pad_height = 256 - spectrogram.shape[0] - spectrogram = np.pad(spectrogram, ((0, pad_height), (0, 0)), mode='constant') - -# Tile the spectrogram -tiles, metadata = tile_spectrogram(spectrogram, tile_size=256, overlap=32) - -# Load model -model = load_model('path/to/model.pt', device='auto') - -# Run inference -predictions = batch_inference(model, tiles, batch_size=32, show_progress=True) - -# Stitch back together -full_prediction = stitch_predictions(predictions, metadata, blend_overlap=True) -``` - -### 4. Post-Processing - -```python -from TokEye.processing import ( - apply_threshold, - remove_small_objects, - compute_statistics, - create_overlay, -) - -# Threshold predictions -binary_mask = apply_threshold(full_prediction, threshold=0.5) - -# Remove small objects -cleaned_mask, num_objects = remove_small_objects(binary_mask, min_size=50) - -# Get statistics -stats = compute_statistics(cleaned_mask) -print(f"Found {stats['num_objects']} objects") -print(f"Mean area: {stats['mean_area']:.1f} pixels") -print(f"Coverage: {stats['coverage']*100:.2f}%") - -# Create visualization -overlay = create_overlay( - spectrogram, - cleaned_mask, - mode='hsv', # or 'white', 'bicolor' - alpha=0.6, -) -# overlay is RGB image, shape: (H, W, 3) -``` - -### 5. Caching for Performance - -```python -from TokEye.processing import CacheManager, generate_cache_key - -# Initialize cache -cache = CacheManager(cache_dir='.cache', max_size_mb=1000) - -# Generate key -params = {'n_fft': 1024, 'hop_length': 128} -key = generate_cache_key(signal, params, prefix='stft') - -# Check cache -if cache.exists(key, 'spectrogram'): - spectrogram = cache.load(key, 'spectrogram') - print("Loaded from cache") -else: - spectrogram = compute_stft(signal, **params) - cache.save(key, spectrogram, cache_type='spectrogram') - print("Computed and cached") - -# View cache stats -stats = cache.get_statistics() -print(f"Cache: {stats['num_entries']} entries, {stats['total_size_mb']:.2f} MB") -``` - -## Complete Pipeline Example - -```python -import numpy as np -from TokEye.processing import * - -# 1. Load signal -signal = np.load('plasma_signal.npy') - -# 2. Preprocess -emphasized = apply_preemphasis(signal, alpha=0.97) -spectrogram = compute_stft(emphasized, n_fft=1024, hop_length=128) - -# 3. Pad to tile size -tile_size = 256 -if spectrogram.shape[0] != tile_size: - pad_height = tile_size - spectrogram.shape[0] - spectrogram = np.pad(spectrogram, ((0, pad_height), (0, 0)), mode='constant') - -# 4. Tile -tiles, metadata = tile_spectrogram(spectrogram, tile_size=tile_size) - -# 5. Infer -model = load_model('model.pt', device='auto') -predictions = batch_inference(model, tiles, batch_size=32) -full_prediction = stitch_predictions(predictions, metadata) - -# 6. Post-process -binary_mask = apply_threshold(full_prediction, threshold=0.5) -cleaned_mask, num_objects = remove_small_objects(binary_mask, min_size=50) - -# 7. Visualize -overlay = create_overlay(spectrogram, cleaned_mask, mode='hsv', alpha=0.6) - -# 8. Get results -stats = compute_statistics(cleaned_mask) -print(f"Detected {num_objects} plasma structures") -print(f"Coverage: {stats['coverage']*100:.2f}%") - -# Save results -import matplotlib.pyplot as plt -plt.imsave('detection_overlay.png', overlay) -``` - -## Function Quick Reference - -### Transforms -- `apply_preemphasis(signal, alpha=0.97)` - Enhance high frequencies -- `compute_stft(signal, n_fft=1024, hop_length=128, ...)` - STFT spectrogram -- `compute_wavelet(signal, wavelet='db8', level=9, ...)` - Wavelet decomposition - -### Tiling -- `tile_spectrogram(spec, tile_size, overlap=0)` - Split into tiles -- `stitch_predictions(tiles, metadata, blend_overlap=True)` - Reconstruct - -### Inference -- `load_model(path, device='auto')` - Load TorchScript model -- `batch_inference(model, tiles, batch_size=32)` - Batch process tiles - -### Post-processing -- `apply_threshold(pred, threshold=0.5)` - Binary threshold -- `remove_small_objects(mask, min_size=50)` - Filter by size -- `create_overlay(spec, mask, mode='white', alpha=0.5)` - Visualize -- `compute_statistics(mask, min_size=0)` - Get detection stats - -### Caching -- `generate_cache_key(data, params, prefix='')` - Create cache key -- `CacheManager(cache_dir, max_size_mb=1000)` - Create cache - - `.save(key, data, cache_type)` - Save to cache - - `.load(key, cache_type)` - Load from cache - - `.exists(key, cache_type)` - Check existence - - `.clear(cache_type=None)` - Clear cache - -## Common Patterns - -### Pattern 1: Process Multiple Signals - -```python -import glob -from pathlib import Path - -for signal_path in glob.glob('data/*.npy'): - signal = np.load(signal_path) - - # Process... - spectrogram = compute_stft(apply_preemphasis(signal)) - - # Save - output_path = Path(signal_path).stem + '_spec.npy' - np.save(output_path, spectrogram) -``` - -### Pattern 2: Cached Processing - -```python -def process_with_cache(signal, cache): - key = generate_cache_key(signal, {}, prefix='processed') - - if cache.exists(key, 'inference'): - return cache.load(key, 'inference') - - # Process pipeline... - result = full_pipeline(signal) - - cache.save(key, result, cache_type='inference') - return result -``` - -### Pattern 3: GPU Batch Processing - -```python -# Process multiple spectrograms -all_tiles = [] -all_metadata = [] - -for spec in spectrograms: - tiles, meta = tile_spectrogram(spec, tile_size=256) - all_tiles.extend(tiles) - all_metadata.append(meta) - -# Single batch inference -model = load_model('model.pt', device='cuda') -all_predictions = batch_inference(model, all_tiles, batch_size=64) - -# Stitch each spectrogram separately -start_idx = 0 -for meta in all_metadata: - end_idx = start_idx + meta['num_tiles'] - predictions = all_predictions[start_idx:end_idx] - result = stitch_predictions(predictions, meta) - start_idx = end_idx -``` - -## Troubleshooting - -### Issue: CUDA out of memory -**Solution**: Reduce batch_size in batch_inference() - -### Issue: Spectrogram height doesn't match tile_size -**Solution**: Pad before tiling: -```python -if spec.shape[0] != tile_size: - pad = tile_size - spec.shape[0] - spec = np.pad(spec, ((0, pad), (0, 0)), mode='constant') -``` - -### Issue: Import error for cv2 -**Solution**: Install OpenCV: -```bash -uv pip install opencv-python -``` - -### Issue: Import error for scipy -**Solution**: Install scipy: -```bash -uv pip install scipy -``` - -## Performance Tips - -1. **Use GPU**: Set `device='cuda'` in load_model() -2. **Batch Processing**: Use larger batch_size if GPU memory allows -3. **Enable Caching**: Cache expensive computations (STFT, wavelet) -4. **Overlap**: Use overlap=0 if blending not needed -5. **Warmup**: Use warmup_model() before batch processing - -## For More Information - -- Full API docs: `src/TokEye/processing/README.md` -- Demo script: `examples/processing_demo.py` -- Implementation details: `PROCESSING_SUMMARY.md` diff --git a/src/tokeye/app/processing/README.md b/src/tokeye/app/processing/README.md deleted file mode 100644 index 67a3506..0000000 --- a/src/tokeye/app/processing/README.md +++ /dev/null @@ -1,459 +0,0 @@ -# TokEye Processing Utilities - -Comprehensive signal processing, model inference, and visualization utilities for the TokEye plasma disruption detection system. - -## Overview - -This module provides a complete pipeline for processing plasma signals, from raw time-domain data to visualized predictions: - -1. **Signal Transformations** - Preemphasis, STFT, Wavelet decomposition -2. **Tiling/Stitching** - Split spectrograms for UNet processing -3. **Model Inference** - Batch processing with PyTorch models -4. **Post-processing** - Thresholding, object filtering, visualization -5. **Caching** - LRU-based caching system for performance - -## Module Structure - -``` -processing/ -├── __init__.py # Main exports -├── transforms.py # Signal processing transformations -├── tiling.py # UNet tiling and stitching -├── inference.py # Model loading and batch inference -├── postprocess.py # Visualization and post-processing -├── cache.py # Caching system with LRU eviction -└── README.md # This file -``` - -## Quick Start - -### Basic Signal Processing Pipeline - -```python -import numpy as np -from TokEye.processing import ( - apply_preemphasis, - compute_stft, - compute_wavelet, -) - -# Load or generate signal -signal = np.random.randn(100000) # Example signal - -# Apply preemphasis filter -emphasized = apply_preemphasis(signal, alpha=0.97) - -# Compute STFT spectrogram -spectrogram = compute_stft( - emphasized, - n_fft=1024, - hop_length=128, - window='hann', - clip_dc=True, -) - -# Compute wavelet decomposition -wavelet_coeffs = compute_wavelet( - signal, - wavelet='db8', - level=9, - mode='sym', - order='freq', -) - -print(f"Spectrogram shape: {spectrogram.shape}") -print(f"Wavelet coeffs shape: {wavelet_coeffs.shape}") -``` - -### Complete Inference Pipeline - -```python -import numpy as np -from TokEye.processing import ( - compute_stft, - tile_spectrogram, - load_model, - batch_inference, - stitch_predictions, - apply_threshold, - remove_small_objects, - create_overlay, -) - -# 1. Create spectrogram -signal = np.random.randn(100000) -spectrogram = compute_stft(signal, n_fft=1024, hop_length=128) - -# 2. Tile for UNet processing -tile_size = 256 -tiles, metadata = tile_spectrogram(spectrogram, tile_size=tile_size) -print(f"Created {len(tiles)} tiles") - -# 3. Load model and run inference -model = load_model('model.pt', device='auto') -predictions = batch_inference(model, tiles, batch_size=32) - -# 4. Stitch predictions back together -full_prediction = stitch_predictions(predictions, metadata) - -# 5. Post-process -binary_mask = apply_threshold(full_prediction, threshold=0.5) -cleaned_mask, num_objects = remove_small_objects(binary_mask, min_size=50) - -# 6. Create visualization -overlay = create_overlay( - spectrogram, - cleaned_mask, - mode='white', - alpha=0.5, -) - -print(f"Detected {num_objects} objects") -``` - -### Using the Cache System - -```python -from TokEye.processing import CacheManager, generate_cache_key -import numpy as np - -# Initialize cache manager -cache = CacheManager( - cache_dir='.cache', - max_size_mb=1000, - max_entries=1000, -) - -# Generate cache key -signal = np.random.randn(10000) -params = {'n_fft': 1024, 'hop_length': 128} -key = generate_cache_key(signal, params, prefix='stft') - -# Check if cached -if cache.exists(key, 'spectrogram'): - spectrogram = cache.load(key, 'spectrogram') - print("Loaded from cache") -else: - # Compute and cache - from TokEye.processing import compute_stft - spectrogram = compute_stft(signal, **params) - cache.save(key, spectrogram, cache_type='spectrogram') - print("Computed and cached") - -# Get cache statistics -stats = cache.get_statistics() -print(f"Cache size: {stats['total_size_mb']:.2f} MB") -print(f"Cache entries: {stats['num_entries']}") -``` - -## API Reference - -### transforms.py - -#### `apply_preemphasis(signal, alpha=0.97)` -Apply preemphasis filter to enhance high frequencies. - -**Parameters:** -- `signal`: Input signal (1D or 2D array) -- `alpha`: Preemphasis coefficient (0-1) - -**Returns:** Preemphasized signal - ---- - -#### `compute_stft(signal, n_fft=1024, hop_length=128, window='hann', clip_dc=True, fs=1.0)` -Compute normalized STFT spectrogram. - -**Parameters:** -- `signal`: Input time-domain signal (1D) -- `n_fft`: FFT size -- `hop_length`: Hop length between frames -- `window`: Window function name -- `clip_dc`: Remove DC component -- `fs`: Sampling frequency - -**Returns:** Normalized STFT magnitude spectrogram (freq_bins, time_frames) - -**Processing Pipeline:** -1. Compute STFT -2. Take magnitude -3. Apply log1p compression -4. Optional DC clipping -5. Mean-std normalization - ---- - -#### `compute_wavelet(signal, wavelet='db8', level=9, mode='sym', order='freq')` -Compute wavelet packet decomposition. - -**Parameters:** -- `signal`: Input signal (1D) -- `wavelet`: Wavelet name ('db8', 'db4', 'haar', etc.) -- `level`: Decomposition level -- `mode`: Signal extension mode -- `order`: Node ordering ('natural' or 'freq') - -**Returns:** 2D array of wavelet coefficients (2^level, coeffs_per_node) - -**Implementation:** -```python -wp = pywt.WaveletPacket(signal, wavelet, mode, maxlevel=level) -nodes = wp.get_level(level, order=order) -values = np.array([n.data for n in nodes], 'd') -values = np.log1p(np.abs(values)) -``` - ---- - -### tiling.py - -#### `tile_spectrogram(spectrogram, tile_size, overlap=0)` -Split spectrogram into square tiles. - -**Parameters:** -- `spectrogram`: Input spectrogram (H, W) or (C, H, W) -- `tile_size`: Size of square tiles (must match height) -- `overlap`: Overlap between tiles in pixels - -**Returns:** Tuple of (tiles_list, metadata_dict) - -**Metadata Keys:** -- `original_width`, `original_height`: Original dimensions -- `tile_size`, `overlap`, `stride`: Tiling parameters -- `num_tiles`: Number of tiles created -- `padding`: Padding added to last tile -- `has_channels`, `num_channels`: Channel information - ---- - -#### `stitch_predictions(tiles, metadata, blend_overlap=True)` -Reconstruct full prediction from tiles. - -**Parameters:** -- `tiles`: List of prediction tiles -- `metadata`: Metadata from tile_spectrogram() -- `blend_overlap`: Average overlapping regions - -**Returns:** Reconstructed full array with padding removed - ---- - -### inference.py - -#### `load_model(model_path, device='auto', map_location=None)` -Load TorchScript model for inference. - -**Parameters:** -- `model_path`: Path to .pt model file -- `device`: Target device ('auto', 'cuda', 'cpu', 'cuda:0', etc.) -- `map_location`: Optional map location for loading - -**Returns:** Loaded model in eval mode - ---- - -#### `batch_inference(model, tiles, batch_size=32, device=None, show_progress=False)` -Run batch inference on tiles. - -**Parameters:** -- `model`: PyTorch model -- `tiles`: List of input tiles (numpy arrays) -- `batch_size`: Batch size for processing -- `device`: Device for inference (None = use model device) -- `show_progress`: Print progress information - -**Returns:** List of predictions (numpy arrays) - ---- - -### postprocess.py - -#### `apply_threshold(prediction, threshold=0.5, binary=True)` -Apply threshold to prediction. - -**Parameters:** -- `prediction`: Input prediction array -- `threshold`: Threshold value -- `binary`: Return binary mask (0/1) vs thresholded values - -**Returns:** Thresholded array - ---- - -#### `remove_small_objects(mask, min_size=50, connectivity=8)` -Remove small connected components. - -**Parameters:** -- `mask`: Binary mask -- `min_size`: Minimum object size in pixels -- `connectivity`: Connectivity (4 or 8) - -**Returns:** Tuple of (cleaned_mask, num_objects) - -**Uses:** OpenCV's `connectedComponentsWithStats` - ---- - -#### `create_overlay(spectrogram, mask, mode='white', alpha=0.5, coherent_color=(0,0,255), transient_color=(0,255,0))` -Create visualization overlay. - -**Parameters:** -- `spectrogram`: Base spectrogram (2D) -- `mask`: Binary or labeled mask -- `mode`: Visualization mode - - `'white'`: Simple white overlay - - `'bicolor'`: Blue=coherent, green=transient - - `'hsv'`: Unique colors per component -- `alpha`: Overlay transparency (0-1) -- `coherent_color`, `transient_color`: BGR colors for bicolor mode - -**Returns:** RGB image (H, W, 3) uint8 - ---- - -#### `compute_statistics(mask, min_size=0)` -Compute object statistics. - -**Parameters:** -- `mask`: Binary mask -- `min_size`: Minimum object size to include - -**Returns:** Dictionary with statistics: -- `num_objects`: Total object count -- `total_area`, `mean_area`, `median_area`, `min_area`, `max_area` -- `coverage`: Fraction of image covered - ---- - -### cache.py - -#### `generate_cache_key(data, params, prefix='')` -Generate unique cache key. - -**Parameters:** -- `data`: Input numpy array -- `params`: Parameters dictionary -- `prefix`: Optional key prefix - -**Returns:** Unique cache key string - -**Method:** SHA256 hash of data + parameters - ---- - -#### `CacheManager(cache_dir='.cache', max_size_mb=1000, max_entries=1000, enable_compression=True)` -Cache manager with LRU eviction. - -**Methods:** - -- `save(key, data, cache_type='general')`: Save data to cache -- `load(key, cache_type='general')`: Load data from cache -- `exists(key, cache_type='general')`: Check if entry exists -- `delete(key, cache_type='general')`: Delete cache entry -- `clear(cache_type=None)`: Clear cache (all or by type) -- `get_statistics()`: Get cache statistics - -**Cache Types:** -- `'spectrogram'`: STFT/wavelet spectrograms -- `'inference'`: Model predictions -- `'wavelet'`: Wavelet decompositions -- `'general'`: General purpose cache - ---- - -## Implementation Notes - -### Signal Processing - -1. **Preemphasis Filter** - - Implements: y[n] = x[n] - α * x[n-1] - - Typical α: 0.95-0.97 - - Enhances high-frequency components - -2. **STFT Pipeline** - - Uses scipy.signal.stft for computation - - Returns: magnitude → log1p → DC clip → normalize - - Normalization: (x - mean) / std - -3. **Wavelet Transform** - - Strict implementation per TokEye spec - - Uses PyWavelets WaveletPacket - - Returns log-compressed coefficients - -### Tiling Strategy - -- **Square Tiles**: Width is divided into height-sized squares -- **Padding**: Last tile padded if width not evenly divisible -- **Overlap**: Optional overlap with averaging during stitching -- **Metadata**: Comprehensive metadata for exact reconstruction - -### Inference Optimization - -- **Batch Processing**: Tiles processed in configurable batches -- **Device Management**: Auto-detect CUDA availability -- **Memory Efficiency**: Uses torch.no_grad() context -- **Model Warmup**: Optional warmup for CUDA kernel compilation - -### Post-processing - -- **Connected Components**: Uses OpenCV for analysis -- **Object Filtering**: Size-based filtering with statistics -- **Visualization Modes**: - - White: Simple binary overlay - - Bicolor: Classify by size heuristic - - HSV: Unique color per component - -### Caching System - -- **LRU Eviction**: Least recently used entries removed first -- **Size Limits**: Both total size (MB) and entry count -- **Type Separation**: Different caches for different data types -- **Persistence**: Metadata saved to disk for recovery -- **Compression**: Optional pickle compression - -## Error Handling - -All functions include: -- Input validation with informative error messages -- Type checking for arrays and parameters -- Graceful handling of edge cases -- Warnings for potential issues -- Import guards for optional dependencies - -## Performance Considerations - -1. **STFT Computation**: O(N log N) for FFT -2. **Wavelet Transform**: O(N) for WaveletPacket -3. **Tiling**: O(N) memory allocation -4. **Inference**: GPU batch processing for speed -5. **Caching**: Constant-time lookup with LRU overhead - -## Dependencies - -- **Required**: numpy, pywavelets, torch -- **Optional**: scipy (STFT), opencv-python (post-processing) - -Install all dependencies: -```bash -uv pip install scipy opencv-python -``` - -## Testing - -Each module includes extensive docstrings with examples. Recommended testing: - -1. **Unit Tests**: Test each function independently -2. **Integration Tests**: Test complete pipeline -3. **Edge Cases**: Empty arrays, single-element arrays, extreme parameters -4. **Performance Tests**: Benchmark critical operations -5. **Round-trip Tests**: Validate tiling/stitching reconstruction - -## Future Enhancements - -- [ ] GPU-accelerated STFT using PyTorch -- [ ] Parallel processing for tiling -- [ ] Advanced visualization modes -- [ ] Configurable cache eviction policies -- [ ] Distributed caching support -- [ ] Model ensemble inference -- [ ] Real-time streaming support diff --git a/src/tokeye/app/processing/__init__.py b/src/tokeye/app/processing/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/tokeye/app/processing/inference.py b/src/tokeye/app/processing/inference.py deleted file mode 100644 index 2324099..0000000 --- a/src/tokeye/app/processing/inference.py +++ /dev/null @@ -1,66 +0,0 @@ -""" -Model Inference Utilities - -This module provides utilities for loading PyTorch models and running -batch inference on tiled spectrograms. -""" - -from pathlib import Path -from venv import logger - -import torch -import torch.nn as nn -from tqdm.auto import tqdm - - -def load_model( - model_path: str | Path, - device: str = "auto", -) -> nn.Module | torch.export.ExportedProgram: - model_path = Path(model_path) - - if not model_path.exists(): - raise FileNotFoundError(f"Model not found: {model_path}") - - if device == "auto": - device = "cuda" if torch.cuda.is_available() else "cpu" - - try: - if model_path.suffix == ".pt2": - module = torch.export.load(str(model_path)) - model = module.module() - else: - model = torch.jit.load( - str(model_path), - map_location=device, - ) - model.eval() - except Exception as e: - raise RuntimeError(f"Failed to load model: {e}") from e - - return model - - -def get_model_info(model: nn.Module) -> dict: - device = next(model.parameters()).device - dtype = next(model.parameters()).dtype - nparams = sum(p.numel() for p in model.parameters()) - - return { - "device": str(device) if device else "unknown", - "dtype": str(dtype) if dtype else "unknown", - "nparams": nparams, - } - - -def warmup( - model: nn.Module, - input_shape: tuple, - num_iterations: int = 5, - dtype: torch.dtype = torch.float32, -): - logger.info("Warming up model...") - dummy_input = torch.randn(*input_shape, device=model.device, dtype=dtype) - with torch.no_grad(): - for _ in tqdm(range(num_iterations)): - _ = model(dummy_input) diff --git a/src/tokeye/app/processing/postprocess.py b/src/tokeye/app/processing/postprocess.py deleted file mode 100644 index 895baba..0000000 --- a/src/tokeye/app/processing/postprocess.py +++ /dev/null @@ -1,386 +0,0 @@ -import cv2 -import numpy as np - - -def apply_threshold( - prediction: np.ndarray, - threshold: float = 0.5, - binary: bool = True, -) -> np.ndarray: - """ - Apply threshold to prediction mask. - - Args: - prediction: Input prediction array with values typically in [0, 1] - threshold: Threshold value for binarization - binary: If True, return binary mask (0 or 1), otherwise return thresholded values - - Returns: - Thresholded prediction array - - Raises: - ValueError: If threshold is not in valid range - - Example: - >>> pred = np.random.rand(256, 256) - >>> mask = apply_threshold(pred, threshold=0.5) - >>> print(np.unique(mask)) # [0, 1] - """ - if not 0 <= threshold <= 1: - raise ValueError(f"Threshold must be in [0, 1], got {threshold}") - - if prediction.size == 0: - return prediction.copy() - - if binary: - # Binary thresholding - thresholded = (prediction >= threshold).astype(np.uint8) - else: - # Keep original values above threshold, zero out below - thresholded = np.where(prediction >= threshold, prediction, 0) - - return thresholded - - -def remove_small_objects( - mask: np.ndarray, - min_size: int = 50, - connectivity: int = 8, -) -> tuple[np.ndarray, int]: - """ - Remove small connected components from binary mask. - - Uses OpenCV's connected components analysis to identify and filter - out small objects based on area. - - Args: - mask: Binary input mask (values should be 0 or 1/255) - min_size: Minimum object size in pixels to keep - connectivity: Connectivity for connected components (4 or 8) - - Returns: - Tuple containing: - - Cleaned binary mask with small objects removed - - Number of remaining objects - - Raises: - ValueError: If mask is not 2D - ValueError: If connectivity is not 4 or 8 - - Example: - >>> mask = np.random.randint(0, 2, (256, 256), dtype=np.uint8) - >>> cleaned, num_objects = remove_small_objects(mask, min_size=100) - >>> print(f"Found {num_objects} objects") - """ - if mask.ndim != 2: - raise ValueError(f"Mask must be 2D, got {mask.ndim}D") - - if connectivity not in [4, 8]: - raise ValueError(f"Connectivity must be 4 or 8, got {connectivity}") - - if min_size < 0: - raise ValueError(f"min_size must be non-negative, got {min_size}") - - if mask.size == 0: - return mask.copy(), 0 - - # Ensure mask is binary uint8 - if mask.dtype != np.uint8: - mask = (mask > 0).astype(np.uint8) - - # Ensure mask is 0 or 255 for cv2 - mask = (mask > 0).astype(np.uint8) * 255 - - # Find connected components - num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats( - mask, connectivity=connectivity - ) - - # Stats columns: [left, top, width, height, area] - # Label 0 is background, so start from 1 - areas = stats[1:, cv2.CC_STAT_AREA] - - # Create output mask - cleaned_mask = np.zeros_like(mask, dtype=np.uint8) - - # Keep only objects larger than min_size - num_kept = 0 - for i in range(1, num_labels): - if areas[i - 1] >= min_size: - cleaned_mask[labels == i] = 255 - num_kept += 1 - - # Convert back to binary (0 or 1) - cleaned_mask = (cleaned_mask > 0).astype(np.uint8) - - return cleaned_mask, num_kept - - -def create_overlay( - spectrogram: np.ndarray, - mask: np.ndarray, - mode: str = "white", - alpha: float = 0.5, - coherent_color: tuple[int, int, int] = (0, 0, 255), # Blue in BGR - transient_color: tuple[int, int, int] = (0, 255, 0), # Green in BGR -) -> np.ndarray: - """ - Create visualization overlay of mask on spectrogram. - - Args: - spectrogram: Input spectrogram array (2D) - mask: Binary or labeled mask (2D) - mode: Visualization mode: - - 'white': Simple white overlay - - 'bicolor': Blue for coherent, red for transient (requires labeled mask) - - 'hsv': Unique colors per component using HSV color space - alpha: Transparency of overlay (0=transparent, 1=opaque) - coherent_color: BGR color for coherent structures (mode='bicolor') - transient_color: BGR color for transient structures (mode='bicolor') - - Returns: - RGB image with overlay (uint8 array of shape (H, W, 3)) - - Raises: - ValueError: If shapes don't match - ValueError: If alpha is not in [0, 1] - ValueError: If mode is invalid - - Example: - >>> spec = np.random.randn(256, 256) - >>> mask = np.random.randint(0, 2, (256, 256)) - >>> overlay = create_overlay(spec, mask, mode='white', alpha=0.5) - >>> print(overlay.shape) # (256, 256, 3) - """ - if spectrogram.shape != mask.shape: - raise ValueError( - f"Spectrogram shape {spectrogram.shape} doesn't match mask shape {mask.shape}" - ) - - if not 0 <= alpha <= 1: - raise ValueError(f"Alpha must be in [0, 1], got {alpha}") - - if mode not in ["white", "bicolor", "hsv"]: - raise ValueError(f"Invalid mode '{mode}', must be 'white', 'bicolor', or 'hsv'") - - # Normalize spectrogram to [0, 255] - spec_min, spec_max = spectrogram.min(), spectrogram.max() - if spec_max - spec_min > 1e-10: - spec_normalized = ( - (spectrogram - spec_min) / (spec_max - spec_min) * 255 - ).astype(np.uint8) - else: - spec_normalized = np.zeros_like(spectrogram, dtype=np.uint8) - - # Convert to BGR (grayscale to color) - base_image = cv2.cvtColor(spec_normalized, cv2.COLOR_GRAY2BGR) - - # Create overlay based on mode - if mode == "white": - # Simple white overlay - overlay = base_image.copy() - overlay[mask > 0] = [255, 255, 255] # White in BGR - - elif mode == "bicolor": - # Bicolor overlay (requires some heuristic to distinguish coherent/transient) - # For now, use a simple heuristic: larger components are coherent, smaller are transient - # This would need to be customized based on actual classification - - try: - num_labels, labels = cv2.connectedComponents( - (mask > 0).astype(np.uint8), connectivity=8 - ) - except Exception as e: - raise RuntimeError(f"Connected components analysis failed: {e}") - - overlay = base_image.copy() - - # Simple heuristic: classify by component size - for i in range(1, num_labels): - component_mask = labels == i - component_size = np.sum(component_mask) - - # Threshold for classification (could be a parameter) - if component_size > 100: # Coherent (larger structures) - overlay[component_mask] = coherent_color - else: # Transient (smaller structures) - overlay[component_mask] = transient_color - - elif mode == "hsv": - # HSV mode: assign unique color to each component - try: - num_labels, labels = cv2.connectedComponents( - (mask > 0).astype(np.uint8), connectivity=8 - ) - except Exception as e: - raise RuntimeError(f"Connected components analysis failed: {e}") - - # Create HSV image - hsv_image = np.zeros((*spectrogram.shape, 3), dtype=np.uint8) - hsv_image[:, :, 1] = 255 # Full saturation - hsv_image[:, :, 2] = 255 # Full value - - # Assign hue based on component label - for i in range(1, num_labels): - component_mask = labels == i - # Distribute hues evenly across spectrum - hue = int((i - 1) * 179 / max(1, num_labels - 1)) - hsv_image[component_mask, 0] = hue - - # Convert HSV to BGR - overlay_color = cv2.cvtColor(hsv_image, cv2.COLOR_HSV2BGR) - - # Blend with base image - overlay = base_image.copy() - mask_bool = mask > 0 - overlay[mask_bool] = overlay_color[mask_bool] - - # Blend overlay with base image using alpha - result = cv2.addWeighted(base_image, 1 - alpha, overlay, alpha, 0) - - # Convert BGR to RGB for standard display - return cv2.cvtColor(result, cv2.COLOR_BGR2RGB) - - - -def compute_channel_threshold_bounds( - prediction: np.ndarray, - num_steps: int = 100, -) -> tuple[float, float]: - """ - Compute lower and upper threshold bounds for a single channel prediction. - - Lower bound: Highest threshold value where applying it results in no pixels (empty mask) - Upper bound: Lowest threshold value where applying it covers the full predicted object - - Args: - prediction: Single channel prediction array (2D) with values in [0, 1] - num_steps: Number of threshold steps to test - - Returns: - Tuple of (lower_bound, upper_bound) normalized to [0, 1] - - Example: - >>> pred = np.random.rand(256, 256) - >>> lower, upper = compute_channel_threshold_bounds(pred) - >>> print(f"Threshold range: [{lower:.3f}, {upper:.3f}]") - """ - if prediction.ndim != 2: - raise ValueError(f"Prediction must be 2D, got {prediction.ndim}D") - - if prediction.size == 0: - return 0.0, 1.0 - - # Get the range of values in the prediction - pred_min = float(prediction.min()) - pred_max = float(prediction.max()) - - if pred_max - pred_min < 1e-10: - # Uniform prediction, return middle value - return pred_min, pred_max - - # Test threshold values from min to max - thresholds = np.linspace(pred_min, pred_max, num_steps) - - # Find lower bound: highest threshold with no pixels - lower_bound = pred_min - for thresh in thresholds: - mask = prediction >= thresh - if np.any(mask): - # Found first threshold that produces pixels - break - lower_bound = thresh - - # Find upper bound: lowest threshold that covers full object - # We define "full object" as the pixels that would be detected at minimum threshold - full_object_mask = ( - prediction >= pred_min + (pred_max - pred_min) * 0.01 - ) # 1% above minimum - - upper_bound = pred_max - for thresh in reversed(thresholds): - mask = prediction >= thresh - # Check if this threshold covers the full object - if np.array_equal(mask, full_object_mask) or np.all(mask[full_object_mask]): - upper_bound = thresh - else: - break - - # Normalize to [0, 1] range - if pred_max - pred_min > 1e-10: - lower_norm = (lower_bound - pred_min) / (pred_max - pred_min) - upper_norm = (upper_bound - pred_min) / (pred_max - pred_min) - else: - lower_norm = 0.0 - upper_norm = 1.0 - - return lower_norm, upper_norm - - -def compute_statistics( - mask: np.ndarray, - min_size: int = 0, -) -> dict: - """ - Compute statistics about detected objects in mask. - - Args: - mask: Binary mask (2D array) - min_size: Minimum object size to include in statistics - - Returns: - Dictionary containing: - - 'num_objects': Total number of objects - - 'total_area': Total area of all objects (pixels) - - 'mean_area': Mean object area - - 'median_area': Median object area - - 'min_area': Minimum object area - - 'max_area': Maximum object area - - 'coverage': Fraction of image covered by objects - - Example: - >>> mask = np.random.randint(0, 2, (256, 256)) - >>> stats = compute_statistics(mask, min_size=50) - >>> print(f"Found {stats['num_objects']} objects") - """ - if mask.ndim != 2: - raise ValueError(f"Mask must be 2D, got {mask.ndim}D") - - # Ensure mask is binary uint8 - mask_binary = (mask > 0).astype(np.uint8) * 255 - - # Find connected components - num_labels, labels, stats_array, centroids = cv2.connectedComponentsWithStats( - mask_binary, connectivity=8 - ) - - # Extract areas (skip background label 0) - areas = stats_array[1:, cv2.CC_STAT_AREA] - - # Filter by minimum size - if min_size > 0: - areas = areas[areas >= min_size] - - # Compute statistics - if len(areas) > 0: - statistics = { - "num_objects": len(areas), - "total_area": int(np.sum(areas)), - "mean_area": float(np.mean(areas)), - "median_area": float(np.median(areas)), - "min_area": int(np.min(areas)), - "max_area": int(np.max(areas)), - "coverage": float(np.sum(areas) / (mask.shape[0] * mask.shape[1])), - } - else: - statistics = { - "num_objects": 0, - "total_area": 0, - "mean_area": 0.0, - "median_area": 0.0, - "min_area": 0, - "max_area": 0, - "coverage": 0.0, - } - - return statistics diff --git a/src/tokeye/app/processing/tiling.py b/src/tokeye/app/processing/tiling.py deleted file mode 100644 index 22ed2ad..0000000 --- a/src/tokeye/app/processing/tiling.py +++ /dev/null @@ -1,292 +0,0 @@ -""" -UNet Tiling and Stitching Utilities - -This module provides functions for splitting spectrograms into tiles for -UNet processing and reconstructing the full prediction from tiled outputs. -""" - -import warnings - -import numpy as np -from TokEye.exceptions import InvalidSpectrogramError, TilingError - - -def tile_spectrogram( - spectrogram: np.ndarray, - tile_size: int, - overlap: int = 0, -) -> tuple[list[np.ndarray], dict]: - """ - Cut spectrogram into square tiles for UNet processing. - - The spectrogram is divided along the width (time axis) into height-sized - squares. If the width is not evenly divisible, the last tile is padded - to match tile_size. - - Args: - spectrogram: Input spectrogram array of shape (height, width) or (channels, height, width) - tile_size: Size of square tiles (must match spectrogram height) - overlap: Number of pixels to overlap between adjacent tiles (default: 0) - - Returns: - Tuple containing: - - List of tile arrays, each of shape (height, tile_size) or (channels, height, tile_size) - - Metadata dict with keys: - - 'original_width': Original spectrogram width - - 'original_height': Original spectrogram height - - 'tile_size': Size of each tile - - 'overlap': Overlap between tiles - - 'num_tiles': Number of tiles created - - 'padding': Amount of padding added to last tile - - 'has_channels': Whether input has channel dimension - - 'num_channels': Number of channels (if has_channels) - - Raises: - ValueError: If spectrogram is not 2D or 3D - ValueError: If tile_size doesn't match height - ValueError: If overlap is negative or >= tile_size - - Example: - >>> spec = np.random.randn(256, 1000) - >>> tiles, metadata = tile_spectrogram(spec, tile_size=256) - >>> print(len(tiles)) # 4 - >>> print(tiles[0].shape) # (256, 256) - """ - if spectrogram.ndim not in [2, 3]: - raise InvalidSpectrogramError( - f"Spectrogram must be 2D (H, W) or 3D (C, H, W), got {spectrogram.ndim}D" - ) - - if overlap < 0: - raise TilingError(f"Overlap must be non-negative, got {overlap}") - - if overlap >= tile_size: - raise TilingError( - f"Overlap ({overlap}) must be less than tile_size ({tile_size})" - ) - - # Handle channel dimension - has_channels = spectrogram.ndim == 3 - if has_channels: - num_channels, height, width = spectrogram.shape - else: - height, width = spectrogram.shape - num_channels = None - - if height != tile_size: - raise TilingError( - f"Spectrogram height ({height}) must match tile_size ({tile_size})" - ) - - if width == 0: - raise TilingError("Spectrogram width must be positive") - - # Calculate stride (step size between tiles) - stride = tile_size - overlap - - # Calculate number of tiles needed - num_tiles = int(np.ceil((width - overlap) / stride)) - - # Calculate total width needed (including padding) - total_width_needed = (num_tiles - 1) * stride + tile_size - padding = total_width_needed - width - - # Pad spectrogram if necessary - if padding > 0: - if has_channels: - pad_width = ((0, 0), (0, 0), (0, padding)) - else: - pad_width = ((0, 0), (0, padding)) - - spectrogram_padded = np.pad( - spectrogram, pad_width, mode="constant", constant_values=0 - ) - else: - spectrogram_padded = spectrogram - - # Extract tiles - tiles = [] - for i in range(num_tiles): - start_idx = i * stride - end_idx = start_idx + tile_size - - if has_channels: - tile = spectrogram_padded[:, :, start_idx:end_idx] - else: - tile = spectrogram_padded[:, start_idx:end_idx] - - tiles.append(tile.copy()) - - # Create metadata for stitching - metadata = { - "original_width": width, - "original_height": height, - "tile_size": tile_size, - "overlap": overlap, - "num_tiles": num_tiles, - "padding": padding, - "stride": stride, - "has_channels": has_channels, - "num_channels": num_channels, - } - - return tiles, metadata - - -def stitch_predictions( - tiles: list[np.ndarray], - metadata: dict, - blend_overlap: bool = True, -) -> np.ndarray: - """ - Reconstruct full prediction from tiles with optional overlap blending. - - Args: - tiles: List of prediction tiles, each of shape matching tile_size from metadata - metadata: Metadata dict returned from tile_spectrogram() - blend_overlap: If True and overlap > 0, blend overlapping regions using averaging - - Returns: - Reconstructed full prediction array with padding removed - - Raises: - ValueError: If tiles list is empty - ValueError: If number of tiles doesn't match metadata - ValueError: If tile shapes are inconsistent - - Example: - >>> spec = np.random.randn(256, 1000) - >>> tiles, metadata = tile_spectrogram(spec, tile_size=256) - >>> # Process tiles through model... - >>> predictions = [model(tile) for tile in tiles] - >>> full_prediction = stitch_predictions(predictions, metadata) - >>> print(full_prediction.shape) # (256, 1000) - """ - if not tiles: - raise TilingError("Tiles list cannot be empty") - - if len(tiles) != metadata["num_tiles"]: - raise TilingError( - f"Number of tiles ({len(tiles)}) doesn't match metadata ({metadata['num_tiles']})" - ) - - # Extract metadata - tile_size = metadata["tile_size"] - overlap = metadata["overlap"] - stride = metadata["stride"] - original_width = metadata["original_width"] - original_height = metadata["original_height"] - has_channels = metadata["has_channels"] - num_channels = metadata["num_channels"] - padding = metadata["padding"] - - # Validate tile shapes - expected_shape = ( - (num_channels, tile_size, tile_size) if has_channels else (tile_size, tile_size) - ) - - for i, tile in enumerate(tiles): - if tile.shape != expected_shape: - raise TilingError( - f"Tile {i} has shape {tile.shape}, expected {expected_shape}" - ) - - # Calculate total reconstructed width (including padding) - reconstructed_width = (len(tiles) - 1) * stride + tile_size - - # Initialize output array - if has_channels: - output_shape = (num_channels, original_height, reconstructed_width) - else: - output_shape = (original_height, reconstructed_width) - - output = np.zeros(output_shape, dtype=tiles[0].dtype) - - # Handle overlap blending - if overlap > 0 and blend_overlap: - # Count how many tiles contribute to each pixel (for averaging) - weight_map = np.zeros(output_shape, dtype=np.float32) - - # Stitch tiles with averaging in overlapping regions - for i, tile in enumerate(tiles): - start_idx = i * stride - end_idx = start_idx + tile_size - - if has_channels: - output[:, :, start_idx:end_idx] += tile - weight_map[:, :, start_idx:end_idx] += 1.0 - else: - output[:, start_idx:end_idx] += tile - weight_map[:, start_idx:end_idx] += 1.0 - - # Normalize by weight map (average overlapping regions) - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", category=RuntimeWarning) - output = np.divide(output, weight_map, where=weight_map > 0) - - else: - # Simple stitching without blending (last tile wins in overlap) - for i, tile in enumerate(tiles): - start_idx = i * stride - end_idx = start_idx + tile_size - - if has_channels: - output[:, :, start_idx:end_idx] = tile - else: - output[:, start_idx:end_idx] = tile - - # Remove padding - if padding > 0: - output = output[:, :, :-padding] if has_channels else output[:, :-padding] - - # Verify final shape matches original - expected_final_shape = ( - (num_channels, original_height, original_width) - if has_channels - else (original_height, original_width) - ) - - if output.shape != expected_final_shape: - warnings.warn( - f"Stitched output shape {output.shape} doesn't match expected " - f"{expected_final_shape}. This may indicate an issue with tiling.", - RuntimeWarning, stacklevel=2, - ) - - return output - - -def validate_tiling_roundtrip( - spectrogram: np.ndarray, - tile_size: int, - overlap: int = 0, - tolerance: float = 1e-6, -) -> bool: - """ - Validate that tiling and stitching correctly reconstruct the input. - - This is a utility function for testing the tiling/stitching pipeline. - - Args: - spectrogram: Input spectrogram to test - tile_size: Tile size for tiling - overlap: Overlap between tiles - tolerance: Maximum allowed absolute difference - - Returns: - True if reconstruction matches input within tolerance - - Example: - >>> spec = np.random.randn(256, 1000) - >>> assert validate_tiling_roundtrip(spec, tile_size=256) - """ - # Tile the spectrogram - tiles, metadata = tile_spectrogram(spectrogram, tile_size, overlap) - - # Stitch back together - reconstructed = stitch_predictions(tiles, metadata, blend_overlap=True) - - # Check if reconstruction matches original - max_diff = np.max(np.abs(spectrogram - reconstructed)) - - return max_diff < tolerance diff --git a/src/tokeye/app/utils/analyze.py b/src/tokeye/app/utils/analyze.py deleted file mode 100644 index 4c73ce4..0000000 --- a/src/tokeye/app/utils/analyze.py +++ /dev/null @@ -1,391 +0,0 @@ -import logging -from pathlib import Path - -import matplotlib.pyplot as plt -import numpy as np -import torch -import torch.nn as nn -from PIL import Image -from skimage import measure -from TokEye.processing.postprocess import apply_threshold, remove_small_objects -from TokEye.processing.transforms import compute_stft -from tqdm.auto import tqdm - -logger = logging.getLogger(__name__) - - -# Signal Functions -def signal_load(file=None, dropdown_path: str = "") -> np.ndarray | None: - filepath = None - if dropdown_path: - filepath = dropdown_path - elif file is not None: - filepath = file.name - else: - logger.error("No file selected or uploaded") - return None - - try: - signal = np.load(filepath) - except Exception as e: - logger.error(f"Failed to load signal: {e}") - return None - if signal.ndim != 1: - logger.error("Signal must be 1D array") - return None - if signal.size == 0: - logger.error("Signal is empty") - return None - return signal - - -# Model Functions -def model_load( - model_path: str | Path, - device: str = "auto", -) -> nn.Module | torch.export.ExportedProgram: - model_path = Path(model_path) - - if not model_path.exists(): - raise FileNotFoundError(f"Model not found: {model_path}") - - if device == "auto": - device = "cuda" if torch.cuda.is_available() else "cpu" - - try: - if model_path.suffix == ".pt2": - module = torch.export.load(str(model_path)) - model = module.module() - else: - model = torch.jit.load( - str(model_path), - map_location=device, - ) - model.eval() - except Exception as e: - raise RuntimeError(f"Failed to load model: {e}") from e - - return model - - -def model_info(model: nn.Module) -> dict: - device = next(model.parameters()).device - dtype = next(model.parameters()).dtype - nparams = sum(p.numel() for p in model.parameters()) - - return { - "device": str(device) if device else "unknown", - "dtype": str(dtype) if dtype else "unknown", - "nparams": nparams, - } - - -def model_warmup( - model: nn.Module, - input_shape: tuple, - num_iterations: int = 5, - dtype: torch.dtype = torch.float32, -): - logger.info("Warming up model...") - device = next(model.parameters()).device - dummy_input = torch.randn(*input_shape, device=device, dtype=dtype) - with torch.no_grad(): - for _ in tqdm(range(num_iterations)): - _ = model(dummy_input) - logger.info("Model warmup complete.") - - -def model_infer( - inp: np.ndarray | None, - model: nn.Module, - mode: str = "single", -) -> np.ndarray | None: - out = None - if mode == "single": - out = torch.from_numpy(inp) - out = out.unsqueeze(0).unsqueeze(0) - out = model(out) - out = out.squeeze(0).squeeze(0) - out = out.numpy() - else: - logger.error(f"Unsupported inference mode: {mode}") - return out - - -# Directory Scanning -def get_available_models() -> list[str]: - """Scan model/ for .pt and .pt2 files.""" - model_dir = Path("model") - if not model_dir.exists(): - return [] - models = list(model_dir.glob("*.pt")) + list(model_dir.glob("*.pt2")) - return [str(m) for m in models] if models else [] - - -def get_available_signals() -> list[str]: - """Scan data/ for .npy files.""" - data_dir = Path("data") - if not data_dir.exists(): - return [] - signals = list(data_dir.glob("*.npy")) - return [str(s) for s in signals] if signals else [] - - -# Pipeline Handlers -def handle_load(model_path: str, signal_path: str) -> tuple: - """Load model and signal, warmup model.""" - try: - if not model_path or not signal_path: - return None, None, "Please select both model and signal" - - model = model_load(model_path) - model_warmup(model, input_shape=(1, 1, 512, 512), num_iterations=3) - signal = signal_load(dropdown_path=signal_path) - - if signal is None: - return None, None, "Failed to load signal" - - status = f"Loaded: {Path(model_path).name} | Signal: {len(signal)} samples" - return model, signal, status - except Exception as e: - logger.error(f"Error in handle_load: {e}") - return None, None, f"Error: {str(e)}" - - -def swap_signal(new_path: str) -> tuple: - """Load new signal without reloading model.""" - try: - if not new_path: - return None, "No signal selected" - signal = signal_load(dropdown_path=new_path) - if signal is None: - return None, "Failed to load signal" - status = f"Signal: {len(signal)} samples" - return signal, status - except Exception as e: - logger.error(f"Error in swap_signal: {e}") - return None, f"Error: {str(e)}" - - -def compute_stft_pipeline( - signal: np.ndarray, - n_fft: int, - hop_length: int, - clip_dc: bool, - percentile_low: float, - percentile_high: float, -) -> tuple: - """Compute STFT and generate preview.""" - try: - if signal is None: - return None, None - - spec = compute_stft( - signal, - n_fft=n_fft, - hop_length=hop_length, - clip_dc=clip_dc, - percentile_low=percentile_low, - percentile_high=percentile_high, - ) - - # Generate preview - fig, ax = plt.subplots(figsize=(8, 4)) - im = ax.imshow(spec, aspect="auto", origin="lower", cmap="viridis") - ax.set_xlabel("Time") - ax.set_ylabel("Frequency") - plt.colorbar(im, ax=ax) - plt.tight_layout() - - # Convert to PIL - fig.canvas.draw() - img = Image.frombytes( - "RGB", fig.canvas.get_width_height(), fig.canvas.tostring_rgb() - ) - plt.close(fig) - - return spec, img - except Exception as e: - logger.error(f"Error in compute_stft_pipeline: {e}") - return None, None - - -def run_inference(spectrogram: np.ndarray, model: nn.Module) -> tuple: - """Run inference and apply sigmoid.""" - try: - if spectrogram is None or model is None: - return None, "Spectrogram or model not loaded" - - inp = torch.from_numpy(spectrogram).unsqueeze(0).unsqueeze(0).float() - device = next(model.parameters()).device - inp = inp.to(device) - - with torch.no_grad(): - out = model(inp) - out = torch.sigmoid(out) - - out = out.squeeze(0).cpu().numpy() - - if out.shape[0] != 2: - return None, f"Expected 2 channels, got {out.shape[0]}" - - status = f"Inference complete: {out.shape}" - return out, status - except Exception as e: - logger.error(f"Error in run_inference: {e}") - return None, f"Error: {str(e)}" - - -# Visualization Renderers -def render_original(spectrogram: np.ndarray) -> Image.Image: - """Render original spectrogram.""" - if spectrogram is None: - return None - - fig, ax = plt.subplots(figsize=(8, 4)) - im = ax.imshow(spectrogram, aspect="auto", origin="lower", cmap="viridis") - ax.set_xlabel("Time") - ax.set_ylabel("Frequency") - plt.colorbar(im, ax=ax) - plt.tight_layout() - - fig.canvas.draw() - img = Image.frombytes( - "RGB", fig.canvas.get_width_height(), fig.canvas.tostring_rgb() - ) - plt.close(fig) - return img - - -def render_enhanced( - inference: np.ndarray, - ch0_enabled: bool, - ch1_enabled: bool, - clip_min: float, - clip_max: float, -) -> Image.Image: - """Render enhanced view with channel overlay.""" - if inference is None: - return None - - # Clip values - ch0 = np.clip(inference[0], clip_min, clip_max) - ch1 = np.clip(inference[1], clip_min, clip_max) - - # Normalize to [0, 1] - ch0 = (ch0 - clip_min) / (clip_max - clip_min) if clip_max > clip_min else ch0 - ch1 = (ch1 - clip_min) / (clip_max - clip_min) if clip_max > clip_min else ch1 - - # Create RGB image - rgb = np.zeros((*ch0.shape, 3)) - - if ch0_enabled and ch1_enabled: - rgb[:, :, 1] = ch0 # Green - rgb[:, :, 0] = ch1 # Red - elif ch0_enabled: - rgb[:, :, 1] = ch0 # Green - elif ch1_enabled: - rgb[:, :, 0] = ch1 # Red - - rgb = (rgb * 255).astype(np.uint8) - return Image.fromarray(rgb) - - -def render_mask( - inference: np.ndarray, - ch0_enabled: bool, - ch1_enabled: bool, - threshold: float, -) -> Image.Image: - """Render binary mask view.""" - if inference is None: - return None - - # Apply threshold - mask_ch0 = apply_threshold(inference[0], threshold, binary=True) - mask_ch1 = apply_threshold(inference[1], threshold, binary=True) - - # Create RGB image - rgb = np.zeros((*mask_ch0.shape, 3)) - - if ch0_enabled and ch1_enabled: - rgb[:, :, 1] = mask_ch0 # Green - rgb[:, :, 0] = mask_ch1 # Red - elif ch0_enabled: - rgb[:, :, 1] = mask_ch0 # Green - elif ch1_enabled: - rgb[:, :, 0] = mask_ch1 # Red - - rgb = (rgb * 255).astype(np.uint8) - return Image.fromarray(rgb) - - -def render_labels( - inference: np.ndarray, - ch0_enabled: bool, - ch1_enabled: bool, - threshold: float, - min_size: int, -) -> Image.Image: - """Render labeled components.""" - if inference is None: - return None - - # Apply threshold and remove small objects - mask_ch0 = apply_threshold(inference[0], threshold, binary=True) - mask_ch1 = apply_threshold(inference[1], threshold, binary=True) - - mask_ch0, _ = remove_small_objects(mask_ch0.astype(np.uint8), min_size=min_size) - mask_ch1, _ = remove_small_objects(mask_ch1.astype(np.uint8), min_size=min_size) - - # Label components - labels_ch0 = measure.label(mask_ch0, connectivity=2) - labels_ch1 = measure.label(mask_ch1, connectivity=2) - - # Create RGB image with unique colors - rgb = np.zeros((*mask_ch0.shape, 3)) - - if ch0_enabled and ch1_enabled: - rgb[:, :, 1] = (labels_ch0 > 0).astype(float) # Green - rgb[:, :, 0] = (labels_ch1 > 0).astype(float) # Red - elif ch0_enabled: - rgb[:, :, 1] = (labels_ch0 > 0).astype(float) # Green - elif ch1_enabled: - rgb[:, :, 0] = (labels_ch1 > 0).astype(float) # Red - - rgb = (rgb * 255).astype(np.uint8) - return Image.fromarray(rgb) - - -def update_visualization( - view_mode: str, - spectrogram: np.ndarray, - inference: np.ndarray, - ch0_enh: bool, - ch1_enh: bool, - clip_min: float, - clip_max: float, - ch0_mask: bool, - ch1_mask: bool, - threshold_mask: float, - ch0_labels: bool, - ch1_labels: bool, - threshold_labels: float, - min_size: int, -) -> Image.Image: - """Route to appropriate renderer based on view mode.""" - try: - if view_mode == "Original": - return render_original(spectrogram) - if view_mode == "Enhanced": - return render_enhanced(inference, ch0_enh, ch1_enh, clip_min, clip_max) - if view_mode == "Mask": - return render_mask(inference, ch0_mask, ch1_mask, threshold_mask) - if view_mode == "Labels": - return render_labels( - inference, ch0_labels, ch1_labels, threshold_labels, min_size - ) - return None - except Exception as e: - logger.error(f"Error in update_visualization: {e}") - return None diff --git a/tests/test_app_load.py b/tests/test_app_load.py new file mode 100644 index 0000000..d8ff7a0 --- /dev/null +++ b/tests/test_app_load.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import numpy as np +import pytest +import torch.nn as nn + +from tokeye.app.analyze import load + + +def test_model_load_delegates_to_hub_and_warmup(monkeypatch): + """model_load should just be hub.load_model(...) + inference.warmup(...).""" + stub_model = nn.Conv2d(1, 2, 1) + calls = {} + + def fake_load_model(source, device): + calls["source"] = source + calls["device"] = device + return stub_model + + def fake_warmup(model): + calls["warmed_up"] = model + + monkeypatch.setattr("tokeye.hub.load_model", fake_load_model) + monkeypatch.setattr("tokeye.inference.warmup", fake_warmup) + + result = load.model_load("big_tf_unet", device="cpu") + + assert result is stub_model + assert calls["source"] == "big_tf_unet" + assert calls["device"] == "cpu" + assert calls["warmed_up"] is stub_model + + +def test_find_models_lists_fake_pt_files(tmp_path, monkeypatch): + monkeypatch.setattr(load, "MODEL_DIR", tmp_path) + (tmp_path / "model_a.pt").touch() + (tmp_path / "model_b.pt2").touch() + (tmp_path / "not_a_model.txt").touch() + + models = load.find_models() + + assert sorted(models) == sorted( + [str(tmp_path / "model_a.pt"), str(tmp_path / "model_b.pt2")] + ) + + +def test_load_single_1d_signal_returns_spectrogram(tmp_path): + signal_path = tmp_path / "signal.npy" + np.save(signal_path, np.random.default_rng(0).normal(size=4096)) + + transform_args = { + "n_fft": 256, + "hop_length": 64, + "clip_dc": True, + "percentile_low": 1.0, + "percentile_high": 99.0, + } + spectrogram = load.load_single(signal_path, transform_args) + + assert spectrogram is not None + assert spectrogram.ndim == 2 + + +def test_load_single_2d_signal_returns_none(tmp_path): + signal_path = tmp_path / "signal_2d.npy" + np.save(signal_path, np.zeros((4, 4))) + + transform_args = {"n_fft": 256, "hop_length": 64} + assert load.load_single(signal_path, transform_args) is None + + +@pytest.mark.parametrize("cached_value", [None, "/path/to/cached.pt"]) +def test_is_model_cached(monkeypatch, cached_value): + monkeypatch.setattr( + load, "try_to_load_from_cache", lambda repo_id, filename: cached_value + ) + + assert load.is_model_cached("big_tf_unet") is (cached_value is not None) + + +def test_create_app_constructs_without_error(): + """Blocks construction should be pure UI wiring: no model load, no network.""" + from tokeye.app.__main__ import create_app + + app = create_app() + + assert app is not None From d91020d3772519edc19417179d059a8cc945e380 Mon Sep 17 00:00:00 2001 From: Nathaniel Chen Date: Sat, 4 Jul 2026 19:07:08 -0400 Subject: [PATCH 3/9] fix: gate one-click Analyze on loaded signal; drop unused SIGNAL_DIR gr.Warning does not halt a gradio .then() chain, so the previous warn-first step still let ensure_model trigger a real model download plus warmup on a no-signal click. ensure_model now takes signal_transform and short-circuits (warn + return model state unchanged) when it is None; downstream model_infer/show_image already no-op on None. Adds three ensure_model tests covering skip/passthrough/ cold-load paths, and removes the now-unused SIGNAL_DIR constant from app/analyze/load.py. --- src/tokeye/app/analyze/analyze.py | 24 ++++++++---------- src/tokeye/app/analyze/load.py | 1 - tests/test_app_load.py | 42 +++++++++++++++++++++++++++++++ 3 files changed, 53 insertions(+), 14 deletions(-) diff --git a/src/tokeye/app/analyze/analyze.py b/src/tokeye/app/analyze/analyze.py index 30a01bd..85d8496 100644 --- a/src/tokeye/app/analyze/analyze.py +++ b/src/tokeye/app/analyze/analyze.py @@ -90,17 +90,19 @@ def wrapper_model_load(model_file): return None -def ensure_model(model, model_file): - """Load the model on first use; pass an already-loaded model through.""" - if model is not None: - return model - return wrapper_model_load(model_file) +def ensure_model(model, model_file, signal_transform): + """Load the model on first use; pass an already-loaded model through. - -def warn_if_missing_signal(signal_transform): - """Nudge the user if Analyze is clicked before any signal is loaded.""" + Skips loading entirely (returning the model state unchanged) if no signal + has been loaded yet: gr.Warning does not halt a .then() chain, so this + gate is what prevents a pointless download/warmup on a no-signal click. + """ if signal_transform is None: gr.Warning("Load a signal first (or click Load Example Signal)") + return model + if model is not None: + return model + return wrapper_model_load(model_file) def wrapper_load_single(signal_directory, signal_file, transform_args): @@ -383,12 +385,8 @@ def analyze_tab(): ## One-click Analyze: ensure a model is loaded, run inference, visualize analyze_btn.click( - fn=warn_if_missing_signal, - inputs=[signal_transform], - outputs=[], - ).then( fn=ensure_model, - inputs=[model, model_file], + inputs=[model, model_file, signal_transform], outputs=[model], ).then( fn=model_infer, diff --git a/src/tokeye/app/analyze/load.py b/src/tokeye/app/analyze/load.py index d740fe6..6891cb0 100644 --- a/src/tokeye/app/analyze/load.py +++ b/src/tokeye/app/analyze/load.py @@ -39,7 +39,6 @@ MODEL_EXTENSIONS = [".pt", ".pt2"] MODEL_DIR = Path("model") SIGNAL_EXTENSIONS = [".npy"] -SIGNAL_DIR = Path("data/input") # Directory Scanning diff --git a/tests/test_app_load.py b/tests/test_app_load.py index d8ff7a0..2a60991 100644 --- a/tests/test_app_load.py +++ b/tests/test_app_load.py @@ -78,6 +78,48 @@ def test_is_model_cached(monkeypatch, cached_value): assert load.is_model_cached("big_tf_unet") is (cached_value is not None) +def test_ensure_model_skips_load_when_no_signal(monkeypatch): + """No signal loaded -> warn and return model state unchanged, no load. + + gr.Warning does not halt a gradio .then() chain, so ensure_model itself + must gate the expensive download/warmup on signal presence. + """ + from tokeye.app.analyze import analyze + + def fail_load(model_file): + raise AssertionError("model load must not run without a signal") + + monkeypatch.setattr(analyze, "wrapper_model_load", fail_load) + + with pytest.warns(UserWarning, match="Load a signal first"): + result = analyze.ensure_model(None, "big_tf_unet", None) + + assert result is None + + +def test_ensure_model_passes_loaded_model_through(monkeypatch): + from tokeye.app.analyze import analyze + + def fail_load(model_file): + raise AssertionError("model load must not run when already loaded") + + monkeypatch.setattr(analyze, "wrapper_model_load", fail_load) + loaded = nn.Conv2d(1, 2, 1) + spectrogram = np.zeros((4, 4)) + + assert analyze.ensure_model(loaded, "big_tf_unet", spectrogram) is loaded + + +def test_ensure_model_loads_when_signal_present_and_model_cold(monkeypatch): + from tokeye.app.analyze import analyze + + stub_model = nn.Conv2d(1, 2, 1) + monkeypatch.setattr(analyze, "wrapper_model_load", lambda model_file: stub_model) + spectrogram = np.zeros((4, 4)) + + assert analyze.ensure_model(None, "big_tf_unet", spectrogram) is stub_model + + def test_create_app_constructs_without_error(): """Blocks construction should be pure UI wiring: no model load, no network.""" from tokeye.app.__main__ import create_app From 134eca6d735594d55df4bd4c0a4b916be56bfc23 Mon Sep 17 00:00:00 2001 From: Nathaniel Chen Date: Sat, 4 Jul 2026 19:19:28 -0400 Subject: [PATCH 4/9] feat: add tokeye CLI with headless batch inference Adds tokeye.batch (gradio-free, Agg-backend batch runner: collect_inputs, load_input, save_overlay_png, process_file, run_batch) and tokeye.cli (argparse entry point with app/run/download/example subcommands, heavy imports deferred into handlers so `tokeye --help` is instant and `tokeye run` never touches gradio). Wires up the `tokeye` console script. --- CITATION.cff | 4 +- pyproject.toml | 7 +- src/tokeye/batch.py | 168 ++++++++++++++++++++++++++++++++++++ src/tokeye/cli.py | 203 ++++++++++++++++++++++++++++++++++++++++++++ tests/test_batch.py | 166 ++++++++++++++++++++++++++++++++++++ tests/test_cli.py | 118 +++++++++++++++++++++++++ uv.lock | 2 +- 7 files changed, 663 insertions(+), 5 deletions(-) create mode 100644 src/tokeye/batch.py create mode 100644 src/tokeye/cli.py create mode 100644 tests/test_batch.py create mode 100644 tests/test_cli.py diff --git a/CITATION.cff b/CITATION.cff index cef3e67..3009f33 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -5,8 +5,8 @@ title: >- message: >- If you use this software, please cite it using the metadata from this file. -version: 0.9.5 -date-released: 2026-06-21 +version: 0.10.0 +date-released: 2026-07-04 doi: 10.5281/zenodo.20791110 authors: - given-names: Nathaniel diff --git a/pyproject.toml b/pyproject.toml index 87f5c5e..fc6644c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,8 +4,8 @@ build-backend = "uv_build" [project] name = "tokeye" -version = "0.9.5" -description = "Add your description here" +version = "0.10.0" +description = "Automatic classification and localization of fluctuating signals in spectrograms" readme = "README.md" requires-python = ">=3.13" dependencies = [ @@ -44,6 +44,9 @@ train = [ "pandas", ] +[project.scripts] +tokeye = "tokeye.cli:main" + [tool.pytest.ini_options] minversion = "9.0" addopts = ["-ra", "-q"] diff --git a/src/tokeye/batch.py b/src/tokeye/batch.py new file mode 100644 index 0000000..b33ca15 --- /dev/null +++ b/src/tokeye/batch.py @@ -0,0 +1,168 @@ +"""Headless batch inference: run TokEye over a list of files with no GUI. + +Never imports gradio, so this module (and anything that imports only this +module) is safe to run on HPC login/compute nodes and in CI. ``matplotlib`` +is switched to the non-interactive ``Agg`` backend before ``pyplot`` is +imported, so this stays headless-safe even without a display. +""" + +from __future__ import annotations + +import glob +import logging +from pathlib import Path + +import matplotlib as mpl +import numpy as np +from tqdm.auto import tqdm + +from . import hub +from .inference import model_infer, signal_to_spectrogram +from .transforms import ( + DEFAULT_CLIP_DC, + DEFAULT_CLIP_HIGH, + DEFAULT_CLIP_LOW, + DEFAULT_HOP, + DEFAULT_N_FFT, +) + +mpl.use("Agg") # Must precede the pyplot import below (headless HPC safety). + +import matplotlib.pyplot as plt # noqa: E402 + +logger = logging.getLogger(__name__) + + +def collect_inputs(inputs: list[str]) -> list[Path]: + """Expand a list of files/directories/glob patterns into concrete paths. + + Each item is resolved as: an existing file (kept as-is), an existing + directory (its ``*.npy`` files, sorted), or otherwise a glob pattern + (matches, sorted). Duplicates are dropped, preserving first-seen order. + """ + collected: list[Path] = [] + for item in inputs: + path = Path(item) + if path.is_file(): + found = [path] + elif path.is_dir(): + found = sorted(path.glob("*.npy")) + else: + # glob.glob (not Path.glob) so absolute patterns keep working: + # Path(".").glob() rejects non-relative patterns outright. + found = sorted(Path(match) for match in glob.glob(item)) # noqa: PTH207 + collected.extend(found) + + seen: set[Path] = set() + result: list[Path] = [] + for path in collected: + if path not in seen: + seen.add(path) + result.append(path) + + if not result: + raise ValueError(f"No input files found for: {inputs}") + + return result + + +def load_input(path: Path, stft_kwargs: dict) -> np.ndarray: + """Load a ``.npy`` file as a spectrogram, computing one if it's a signal.""" + arr = np.load(path) + if arr.ndim == 1: + return signal_to_spectrogram(arr, **stft_kwargs) + if arr.ndim == 2: + return arr.astype(float) + raise ValueError(f"expected 1D signal or 2D spectrogram, got ndim={arr.ndim}") + + +def save_overlay_png( + spectrogram: np.ndarray, + mask: np.ndarray, + out_path: Path, + threshold: float = 0.5, + dpi: int = 150, +) -> None: + """Save a grayscale spectrogram with a semi-transparent mask overlay. + + Coherent activity (``mask[0]``) is tinted green, transient activity + (``mask[1]``) is tinted red, both thresholded at ``threshold``. + """ + height, width = spectrogram.shape + overlay = np.zeros((height, width, 4), dtype=np.float32) + overlay[mask[0] >= threshold] = (0.0, 1.0, 0.0, 0.4) + overlay[mask[1] >= threshold] = (1.0, 0.0, 0.0, 0.4) + + fig, ax = plt.subplots() + ax.imshow(spectrogram, cmap="gray", origin="lower", aspect="auto") + ax.imshow(overlay, origin="lower", aspect="auto") + fig.tight_layout() + fig.savefig(out_path, dpi=dpi) + plt.close(fig) + + +def process_file( + path: Path, + model, + stft_kwargs: dict, + out_dir: Path, + save_png: bool = True, + threshold: float = 0.5, +) -> Path: + """Run inference on a single input file and write its outputs to disk.""" + spectrogram = load_input(path, stft_kwargs) + mask = model_infer(spectrogram, model) + + mask_path = out_dir / f"{path.stem}_mask.npy" + np.save(mask_path, mask.astype(np.float32)) + + if save_png: + preview_path = out_dir / f"{path.stem}_preview.png" + save_overlay_png(spectrogram, mask, preview_path, threshold=threshold) + + return mask_path + + +def run_batch( + inputs: list[str], + model: str | Path = hub.DEFAULT_MODEL, + out_dir: Path = Path("tokeye_output"), + stft_kwargs: dict | None = None, + save_png: bool = True, + threshold: float = 0.5, + device: str = "auto", +) -> int: + """Run inference over ``inputs``, writing masks (and previews) to ``out_dir``. + + Returns the number of files that failed to process. + """ + paths = collect_inputs(inputs) + resolved_stft_kwargs = stft_kwargs if stft_kwargs is not None else { + "n_fft": DEFAULT_N_FFT, + "hop": DEFAULT_HOP, + "clip_dc": DEFAULT_CLIP_DC, + "clip_low": DEFAULT_CLIP_LOW, + "clip_high": DEFAULT_CLIP_HIGH, + } + + loaded_model = hub.load_model(model, device) + + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + failures = 0 + for path in tqdm(paths, desc="tokeye run"): + try: + process_file( + path, + loaded_model, + resolved_stft_kwargs, + out_dir, + save_png=save_png, + threshold=threshold, + ) + except Exception: + logger.error("Failed to process %s", path, exc_info=True) + failures += 1 + + return failures diff --git a/src/tokeye/cli.py b/src/tokeye/cli.py new file mode 100644 index 0000000..a8c62ab --- /dev/null +++ b/src/tokeye/cli.py @@ -0,0 +1,203 @@ +"""``tokeye`` console entry point. + +Argparse only (no new dependencies). Heavy imports (torch, ``tokeye.batch``, +``tokeye.app``) are deferred into each subcommand handler so ``tokeye --help`` +returns instantly and ``tokeye run`` never imports gradio. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from typing import TYPE_CHECKING + +from tokeye.transforms import ( + DEFAULT_CLIP_HIGH, + DEFAULT_CLIP_LOW, + DEFAULT_HOP, + DEFAULT_N_FFT, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + + +def _add_app_subcommand(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser("app", help="Launch the TokEye Gradio web app.") + parser.add_argument( + "--port", type=int, default=7860, help="Port to serve the app on." + ) + parser.add_argument( + "--share", action="store_true", help="Create a public Gradio share link." + ) + parser.add_argument( + "--open", + dest="open_browser", + action="store_true", + help="Open the app in a browser on launch.", + ) + parser.set_defaults(handler=_handle_app) + + +def _add_run_subcommand(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser("run", help="Run batch inference on one or more inputs.") + parser.add_argument( + "inputs", + nargs="+", + metavar="INPUT", + help="Files, directories of .npy files, or glob patterns.", + ) + parser.add_argument( + "--model", + default=None, + help="Registry name or path to a model checkpoint (default: big_tf_unet).", + ) + parser.add_argument( + "--output-dir", + default="tokeye_output", + help="Directory to write masks and previews to.", + ) + parser.add_argument("--n-fft", type=int, default=DEFAULT_N_FFT) + parser.add_argument("--hop", type=int, default=DEFAULT_HOP) + parser.add_argument( + "--keep-dc", + action="store_true", + help="Do not clip the DC bin (clipped by default).", + ) + parser.add_argument("--clip-low", type=float, default=DEFAULT_CLIP_LOW) + parser.add_argument("--clip-high", type=float, default=DEFAULT_CLIP_HIGH) + parser.add_argument("--threshold", type=float, default=0.5) + parser.add_argument( + "--no-png", + dest="save_png", + action="store_false", + help="Skip PNG overlay previews.", + ) + parser.add_argument("--device", default="auto") + parser.set_defaults(handler=_handle_run) + + +def _add_download_subcommand(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser("download", help="Download one or more model checkpoints.") + parser.add_argument( + "models", + nargs="*", + default=None, + metavar="MODEL", + help="Model registry name(s) to download (default: big_tf_unet).", + ) + parser.set_defaults(handler=_handle_download) + + +def _add_example_subcommand(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser( + "example", help="Write a synthetic example signal to a .npy file." + ) + parser.add_argument("--output", default="tokeye_example.npy") + parser.add_argument("--duration", type=float, default=2.0) + parser.add_argument("--fs", type=float, default=200_000.0) + parser.add_argument("--seed", type=int, default=0) + parser.set_defaults(handler=_handle_example) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="tokeye", + description=( + "Automatic classification and localization of fluctuating signals " + "in spectrograms." + ), + ) + parser.add_argument( + "--version", action="store_true", help="Print the tokeye version and exit." + ) + subparsers = parser.add_subparsers(dest="command") + _add_app_subcommand(subparsers) + _add_run_subcommand(subparsers) + _add_download_subcommand(subparsers) + _add_example_subcommand(subparsers) + return parser + + +def _handle_app(args: argparse.Namespace) -> int: + from tokeye.app.__main__ import main as app_main + + app_main(port=args.port, share=args.share, open_browser=args.open_browser) + return 0 + + +def _handle_run(args: argparse.Namespace) -> int: + from tokeye import batch + from tokeye.hub import DEFAULT_MODEL + + stft_kwargs = { + "n_fft": args.n_fft, + "hop": args.hop, + "clip_dc": not args.keep_dc, + "clip_low": args.clip_low, + "clip_high": args.clip_high, + } + model = args.model if args.model is not None else DEFAULT_MODEL + + try: + return batch.run_batch( + args.inputs, + model=model, + out_dir=Path(args.output_dir), + stft_kwargs=stft_kwargs, + save_png=args.save_png, + threshold=args.threshold, + device=args.device, + ) + except ValueError as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + + +def _handle_download(args: argparse.Namespace) -> int: + from tokeye.hub import DEFAULT_MODEL, download_model + + names = args.models or [DEFAULT_MODEL] + for name in names: + try: + path = download_model(name) + except ValueError as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + print(path) + return 0 + + +def _handle_example(args: argparse.Namespace) -> int: + import numpy as np + + from tokeye.examples import make_example_signal + + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + signal = make_example_signal(duration_s=args.duration, fs=args.fs, seed=args.seed) + np.save(output_path, signal) + print(output_path) + return 0 + + +def main(argv: Sequence[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + + if args.version: + from importlib.metadata import version + + print(version("tokeye")) + return 0 + + if args.command is None: + parser.print_help() + return 2 + + return args.handler(args) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_batch.py b/tests/test_batch.py new file mode 100644 index 0000000..c6d5936 --- /dev/null +++ b/tests/test_batch.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +import numpy as np +import pytest +import torch.nn as nn + +from tokeye import batch + + +@pytest.fixture +def stub_model(monkeypatch): + """A cheap Conv2d(1, 2, 1) stands in for the real (heavy) TokEye model.""" + model = nn.Conv2d(1, 2, kernel_size=1) + monkeypatch.setattr("tokeye.hub.load_model", lambda source, device: model) + return model + + +@pytest.fixture +def signal_npy(tmp_path): + path = tmp_path / "signal.npy" + sig = np.random.default_rng(0).normal(size=8192).astype(np.float32) + np.save(path, sig) + return path + + +@pytest.fixture +def spectrogram_npy(tmp_path): + path = tmp_path / "spectrogram.npy" + spec = np.random.default_rng(1).normal(size=(64, 32)).astype(np.float32) + np.save(path, spec) + return path + + +class TestCollectInputs: + def test_single_file(self, signal_npy): + assert batch.collect_inputs([str(signal_npy)]) == [signal_npy] + + def test_directory_finds_only_npy_sorted(self, tmp_path): + (tmp_path / "b.npy").touch() + (tmp_path / "a.npy").touch() + (tmp_path / "ignore.txt").touch() + + result = batch.collect_inputs([str(tmp_path)]) + + assert result == [tmp_path / "a.npy", tmp_path / "b.npy"] + + def test_glob_pattern(self, tmp_path): + (tmp_path / "x1.npy").touch() + (tmp_path / "x2.npy").touch() + + result = batch.collect_inputs([str(tmp_path / "x*.npy")]) + + assert result == [tmp_path / "x1.npy", tmp_path / "x2.npy"] + + def test_mixed_inputs_dedup_preserving_order(self, tmp_path): + (tmp_path / "a.npy").touch() + (tmp_path / "b.npy").touch() + + # The directory glob would also match a.npy; it should appear once, + # in its first-seen position. + result = batch.collect_inputs( + [str(tmp_path / "a.npy"), str(tmp_path)] + ) + + assert result == [tmp_path / "a.npy", tmp_path / "b.npy"] + + def test_empty_result_raises_value_error(self, tmp_path): + with pytest.raises(ValueError, match="No input files found"): + batch.collect_inputs([str(tmp_path / "does_not_exist_*.npy")]) + + +class TestLoadInput: + def test_1d_signal_becomes_spectrogram(self, signal_npy): + spec = batch.load_input(signal_npy, {"n_fft": 256, "hop": 64}) + assert spec.ndim == 2 + + def test_2d_spectrogram_passes_through(self, spectrogram_npy): + spec = batch.load_input(spectrogram_npy, {}) + assert spec.shape == (64, 32) + assert np.issubdtype(spec.dtype, np.floating) + + def test_3d_array_raises_value_error(self, tmp_path): + path = tmp_path / "bad.npy" + np.save(path, np.zeros((2, 3, 4))) + + with pytest.raises(ValueError, match="ndim=3"): + batch.load_input(path, {}) + + +class TestRunBatch: + def test_on_1d_signal_writes_mask_and_preview( + self, stub_model, signal_npy, tmp_path + ): + out_dir = tmp_path / "out" + failures = batch.run_batch( + [str(signal_npy)], + out_dir=out_dir, + stft_kwargs={"n_fft": 256, "hop": 64}, + ) + + assert failures == 0 + mask_path = out_dir / "signal_mask.npy" + preview_path = out_dir / "signal_preview.png" + assert mask_path.exists() + assert preview_path.exists() + assert preview_path.stat().st_size > 0 + + mask = np.load(mask_path) + assert mask.ndim == 3 + assert mask.shape[0] == 2 + assert mask.dtype == np.float32 + assert np.all(mask >= 0.0) and np.all(mask <= 1.0) + + def test_on_2d_spectrogram_writes_mask_and_preview( + self, stub_model, spectrogram_npy, tmp_path + ): + out_dir = tmp_path / "out" + failures = batch.run_batch([str(spectrogram_npy)], out_dir=out_dir) + + assert failures == 0 + mask = np.load(out_dir / "spectrogram_mask.npy") + assert mask.shape == (2, 64, 32) + assert mask.dtype == np.float32 + assert np.all(mask >= 0.0) and np.all(mask <= 1.0) + assert (out_dir / "spectrogram_preview.png").exists() + + def test_save_png_false_skips_preview(self, stub_model, spectrogram_npy, tmp_path): + out_dir = tmp_path / "out" + failures = batch.run_batch( + [str(spectrogram_npy)], out_dir=out_dir, save_png=False + ) + + assert failures == 0 + assert (out_dir / "spectrogram_mask.npy").exists() + assert not (out_dir / "spectrogram_preview.png").exists() + + def test_one_bad_file_among_good_ones_counts_as_failure( + self, stub_model, spectrogram_npy, tmp_path + ): + bad_path = tmp_path / "bad.npy" + np.save(bad_path, np.zeros((2, 3, 4))) + + out_dir = tmp_path / "out" + failures = batch.run_batch( + [str(spectrogram_npy), str(bad_path)], out_dir=out_dir + ) + + assert failures == 1 + assert (out_dir / "spectrogram_mask.npy").exists() + assert not (out_dir / "bad_mask.npy").exists() + + def test_loads_model_once_via_hub(self, spectrogram_npy, tmp_path, monkeypatch): + model = nn.Conv2d(1, 2, kernel_size=1) + calls = [] + + def fake_load_model(source, device): + calls.append((source, device)) + return model + + monkeypatch.setattr("tokeye.hub.load_model", fake_load_model) + + batch.run_batch( + [str(spectrogram_npy)], out_dir=tmp_path / "out", device="cpu" + ) + + assert calls == [(batch.hub.DEFAULT_MODEL, "cpu")] diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..3d68baa --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from tokeye.cli import build_parser, main + + +class TestBuildParser: + def test_run_subcommand_parses_options(self): + parser = build_parser() + + args = parser.parse_args( + ["run", "a.npy", "--hop", "128", "--no-png", "--model", "model/x.pt"] + ) + + assert args.inputs == ["a.npy"] + assert args.hop == 128 + assert args.save_png is False + assert args.model == "model/x.pt" + + def test_run_subcommand_defaults(self): + parser = build_parser() + + args = parser.parse_args(["run", "a.npy"]) + + assert args.model is None # resolved to DEFAULT_MODEL in the handler + assert args.output_dir == "tokeye_output" + assert args.n_fft == 1024 + assert args.hop == 256 + assert args.keep_dc is False + assert args.clip_low == 1.0 + assert args.clip_high == 99.0 + assert args.threshold == 0.5 + assert args.save_png is True + assert args.device == "auto" + + def test_app_subcommand_defaults(self): + parser = build_parser() + + args = parser.parse_args(["app"]) + + assert args.port == 7860 + assert args.share is False + assert args.open_browser is False + + +class TestMain: + def test_version_returns_zero_and_prints_version(self, capsys): + exit_code = main(["--version"]) + + assert exit_code == 0 + out = capsys.readouterr().out + assert out.strip() # some version string was printed + + def test_no_subcommand_returns_two_or_exits(self, capsys): + try: + exit_code = main([]) + except SystemExit as exc: + assert exc.code == 2 + else: + assert exit_code == 2 + + def test_run_with_nonexistent_input_returns_two_clean_error(self, capsys): + exit_code = main(["run", "does_not_exist_anywhere_xyz.npy"]) + + assert exit_code == 2 + err = capsys.readouterr().err + assert "does_not_exist_anywhere_xyz.npy" in err + + def test_example_writes_file_and_returns_zero(self, tmp_path, capsys): + out_path = tmp_path / "example.npy" + + exit_code = main( + [ + "example", + "--output", + str(out_path), + "--duration", + "0.01", + "--fs", + "10000", + ] + ) + + assert exit_code == 0 + assert out_path.exists() + sig = np.load(out_path) + assert sig.shape[0] == 100 + + def test_download_unknown_model_returns_two_clean_error(self, capsys): + exit_code = main(["download", "not_a_real_model_name"]) + + assert exit_code == 2 + err = capsys.readouterr().err + assert "not_a_real_model_name" in err + + def test_app_subcommand_delegates_to_app_main(self, monkeypatch): + calls = {} + + def fake_app_main(port, share, open_browser): + calls["port"] = port + calls["share"] = share + calls["open_browser"] = open_browser + + monkeypatch.setattr("tokeye.app.__main__.main", fake_app_main) + + exit_code = main(["app", "--port", "1234", "--share"]) + + assert exit_code == 0 + assert calls == {"port": 1234, "share": True, "open_browser": False} + + +@pytest.mark.parametrize("argv", [["--help"], ["run", "--help"]]) +def test_help_does_not_crash(argv, capsys): + with pytest.raises(SystemExit) as exc_info: + main(argv) + assert exc_info.value.code == 0 diff --git a/uv.lock b/uv.lock index 4b6a0bd..bf09a78 100644 --- a/uv.lock +++ b/uv.lock @@ -2707,7 +2707,7 @@ wheels = [ [[package]] name = "tokeye" -version = "0.9.5" +version = "0.10.0" source = { editable = "." } dependencies = [ { name = "gradio" }, From 5c55b28773da8dbf186ba13d656ec3319134a2ee Mon Sep 17 00:00:00 2001 From: Nathaniel Chen Date: Sat, 4 Jul 2026 19:28:11 -0400 Subject: [PATCH 5/9] fix: CLI review fixes for example output path and run error handling - tokeye example: normalize extension-less --output to .npy before saving so the printed path is the file np.save actually writes - tokeye run: catch FileNotFoundError from a mistyped --model path and emit a clean stderr error (exit 2) instead of a raw traceback - test_loads_model_once_via_hub: use two input files so once-per-run is distinguishable from once-per-file --- src/tokeye/cli.py | 6 +++++- tests/test_batch.py | 18 +++++++++++++++--- tests/test_cli.py | 34 ++++++++++++++++++++++++++++++++++ 3 files changed, 54 insertions(+), 4 deletions(-) diff --git a/src/tokeye/cli.py b/src/tokeye/cli.py index a8c62ab..a490bf7 100644 --- a/src/tokeye/cli.py +++ b/src/tokeye/cli.py @@ -150,7 +150,7 @@ def _handle_run(args: argparse.Namespace) -> int: threshold=args.threshold, device=args.device, ) - except ValueError as exc: + except (ValueError, FileNotFoundError) as exc: print(f"error: {exc}", file=sys.stderr) return 2 @@ -175,6 +175,10 @@ def _handle_example(args: argparse.Namespace) -> int: from tokeye.examples import make_example_signal output_path = Path(args.output) + if output_path.suffix != ".npy": + # np.save silently appends ".npy" to paths without that suffix; + # normalize up-front so the printed path is the file that exists. + output_path = output_path.with_suffix(".npy") output_path.parent.mkdir(parents=True, exist_ok=True) signal = make_example_signal(duration_s=args.duration, fs=args.fs, seed=args.seed) np.save(output_path, signal) diff --git a/tests/test_batch.py b/tests/test_batch.py index c6d5936..7789735 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -149,7 +149,9 @@ def test_one_bad_file_among_good_ones_counts_as_failure( assert (out_dir / "spectrogram_mask.npy").exists() assert not (out_dir / "bad_mask.npy").exists() - def test_loads_model_once_via_hub(self, spectrogram_npy, tmp_path, monkeypatch): + def test_loads_model_once_via_hub(self, tmp_path, monkeypatch): + """Multiple input files -> exactly one load_model call (per run, not + per file).""" model = nn.Conv2d(1, 2, kernel_size=1) calls = [] @@ -159,8 +161,18 @@ def fake_load_model(source, device): monkeypatch.setattr("tokeye.hub.load_model", fake_load_model) - batch.run_batch( - [str(spectrogram_npy)], out_dir=tmp_path / "out", device="cpu" + rng = np.random.default_rng(2) + for name in ("first.npy", "second.npy"): + np.save(tmp_path / name, rng.normal(size=(64, 32)).astype(np.float32)) + + out_dir = tmp_path / "out" + failures = batch.run_batch( + [str(tmp_path / "first.npy"), str(tmp_path / "second.npy")], + out_dir=out_dir, + device="cpu", ) + assert failures == 0 + assert (out_dir / "first_mask.npy").exists() + assert (out_dir / "second_mask.npy").exists() assert calls == [(batch.hub.DEFAULT_MODEL, "cpu")] diff --git a/tests/test_cli.py b/tests/test_cli.py index 3d68baa..2be5bc0 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -68,6 +68,21 @@ def test_run_with_nonexistent_input_returns_two_clean_error(self, capsys): err = capsys.readouterr().err assert "does_not_exist_anywhere_xyz.npy" in err + def test_run_with_missing_model_path_returns_two_clean_error( + self, tmp_path, capsys + ): + input_path = tmp_path / "input.npy" + np.save(input_path, np.zeros((64, 32), dtype=np.float32)) + + exit_code = main( + ["run", str(input_path), "--model", "nope/missing.pt"] + ) + + assert exit_code == 2 + err = capsys.readouterr().err + assert "nope/missing.pt" in err + assert "Traceback" not in err + def test_example_writes_file_and_returns_zero(self, tmp_path, capsys): out_path = tmp_path / "example.npy" @@ -88,6 +103,25 @@ def test_example_writes_file_and_returns_zero(self, tmp_path, capsys): sig = np.load(out_path) assert sig.shape[0] == 100 + def test_example_extensionless_output_prints_actual_file(self, tmp_path, capsys): + """np.save appends .npy; the printed path must be the file that exists.""" + exit_code = main( + [ + "example", + "--output", + str(tmp_path / "demo"), + "--duration", + "0.01", + "--fs", + "10000", + ] + ) + + assert exit_code == 0 + printed = capsys.readouterr().out.strip() + assert printed == str(tmp_path / "demo.npy") + assert (tmp_path / "demo.npy").exists() + def test_download_unknown_model_returns_two_clean_error(self, capsys): exit_code = main(["download", "not_a_real_model_name"]) From 8a9b944eb2b50ada5059a842b70be6b26eb9329c Mon Sep 17 00:00:00 2001 From: Nathaniel Chen Date: Sat, 4 Jul 2026 19:35:12 -0400 Subject: [PATCH 6/9] feat: add maintainer model-upload script with verification gate --- scripts/upload_model.py | 169 +++++++++++++++++++++++++++++++++++++ tests/test_upload_model.py | 50 +++++++++++ 2 files changed, 219 insertions(+) create mode 100644 scripts/upload_model.py create mode 100644 tests/test_upload_model.py diff --git a/scripts/upload_model.py b/scripts/upload_model.py new file mode 100644 index 0000000..2ea3b0f --- /dev/null +++ b/scripts/upload_model.py @@ -0,0 +1,169 @@ +"""Publish a verified TokEye checkpoint to the Hugging Face Hub. + +Run this ONCE per checkpoint update, by a maintainer with write access to the +target Hub repo. `tokeye.hub.download_model` is what every user's install +pulls weights from — this script's verification gate is the only thing +standing between a bad artifact and a broken install for all of them. + +Prereq: authenticate first, either via `hf auth login` or by setting the +`HF_TOKEN` environment variable, with write access to the target org/repo. + +Usage: + + uv run python scripts/upload_model.py --create + uv run python scripts/upload_model.py --model big_tf_unet \ + --repo PlasmaControl/tokeye + +That is equivalent to the raw one-liner: + + hf upload PlasmaControl/tokeye model/big_tf_unet_251210.pt \ + big_tf_unet_251210.pt + +...except that `hf upload` does not check that the file loads into the +architecture `tokeye.hub` expects. This script's added value is the +verification gate in `verify_checkpoint`: it refuses to upload anything that +does not load, strictly, into the registered model class and survive a tiny +forward pass. That is what protects the repo from an accidental upload of a +file like `model/big_tf_unet_251210_weights.pt`, which uses old +(`in_conv.double_conv.*`-style) key names and would silently break every +downstream user. +""" + +from __future__ import annotations + +import argparse +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from huggingface_hub import HfApi +from huggingface_hub.utils import HfHubHTTPError, LocalTokenNotFoundError + +from tokeye.hub import DEFAULT_MODEL, DEFAULT_REPO_ID, MODEL_REGISTRY + +if TYPE_CHECKING: + from collections.abc import Callable + +_PROBE_SHAPE = (1, 1, 64, 64) + + +def verify_checkpoint(path: Path, builder: Callable[[], nn.Module]) -> None: + """Refuse (via ``SystemExit``) to proceed unless ``path`` is a good checkpoint. + + "Good" means: a weights-only state dict that loads strictly into a fresh + instance of the architecture ``builder`` returns, and survives a tiny + forward pass. This is the safety gate the whole script exists for. + """ + try: + state_dict = torch.load(path, map_location="cpu", weights_only=True) + except Exception as exc: + raise SystemExit( + f"error: could not load {path} as a weights-only state dict " + f"({exc}).\nIf this is a pickled full nn.Module (a legacy " + "checkpoint), it must be converted to a plain state_dict before " + "upload. Refusing to upload." + ) from exc + + if not isinstance(state_dict, Mapping): + raise SystemExit( + f"error: {path} loaded as a {type(state_dict).__name__}, not a " + "state_dict mapping. Refusing to upload." + ) + + model = builder() + try: + model.load_state_dict(state_dict, strict=True) + except RuntimeError as exc: + mismatches = "\n".join(str(exc).splitlines()[:8]) + raise SystemExit( + f"error: {path} does not match the {type(model).__name__} " + f"architecture. Refusing to upload.\n{mismatches}" + ) from exc + + model.eval() + with torch.no_grad(): + try: + model(torch.randn(*_PROBE_SHAPE)) + except Exception as exc: + raise SystemExit( + f"error: {path} loaded but failed a forward-pass sanity " + f"check ({exc}). Refusing to upload." + ) from exc + + print(f"verified: {path} matches {type(model).__name__} and runs forward pass.") + + +def _resolve_file(model_name: str, file_arg: str | None) -> Path: + if file_arg is not None: + return Path(file_arg) + return Path("model") / MODEL_REGISTRY[model_name].filename + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="upload_model.py", + description=( + "Verify a TokEye checkpoint against its registered architecture, " + "then publish it to the Hugging Face Hub." + ), + ) + parser.add_argument( + "--model", + default=DEFAULT_MODEL, + choices=sorted(MODEL_REGISTRY), + help="Registry entry to publish (default: %(default)s).", + ) + parser.add_argument( + "--file", + default=None, + help="Local checkpoint to upload (default: model/).", + ) + parser.add_argument( + "--repo", + default=DEFAULT_REPO_ID, + help="Target Hugging Face Hub repo id (default: %(default)s).", + ) + parser.add_argument( + "--create", + action="store_true", + help="Create the repo first (HfApi().create_repo(repo, exist_ok=True)).", + ) + return parser + + +def main() -> int: + args = build_parser().parse_args() + spec = MODEL_REGISTRY[args.model] + file_path = _resolve_file(args.model, args.file) + + if not file_path.exists(): + raise SystemExit(f"error: checkpoint not found: {file_path}") + + verify_checkpoint(file_path, spec.builder) + + api = HfApi() + try: + if args.create: + api.create_repo(args.repo, exist_ok=True) + commit_info = api.upload_file( + path_or_fileobj=file_path, + path_in_repo=spec.filename, + repo_id=args.repo, + commit_message=f"Upload {spec.filename} (tokeye model={args.model!r})", + ) + except (HfHubHTTPError, LocalTokenNotFoundError) as exc: + raise SystemExit( + f"error: Hugging Face Hub upload failed ({exc}).\n" + "Run `hf auth login` (or set the HF_TOKEN env var) with write " + f"access to {args.repo!r} and try again." + ) from exc + + print(f"uploaded: {commit_info.commit_url}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_upload_model.py b/tests/test_upload_model.py new file mode 100644 index 0000000..3484308 --- /dev/null +++ b/tests/test_upload_model.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest +import torch +import torch.nn as nn + +from tokeye.hub import MODEL_REGISTRY + +_SCRIPT_PATH = Path(__file__).resolve().parents[1] / "scripts" / "upload_model.py" + + +def _load_upload_model_module(): + spec = importlib.util.spec_from_file_location("upload_model", _SCRIPT_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +upload_model = _load_upload_model_module() + + +def test_verify_checkpoint_accepts_good_state_dict(tmp_path): + spec = MODEL_REGISTRY["big_tf_unet"] + path = tmp_path / "good.pt" + torch.save(spec.builder().state_dict(), path) + + upload_model.verify_checkpoint(path, spec.builder) # must not raise + + +def test_verify_checkpoint_rejects_renamed_keys(tmp_path): + spec = MODEL_REGISTRY["big_tf_unet"] + path = tmp_path / "renamed.pt" + state_dict = spec.builder().state_dict() + renamed = {f"old.{key}": value for key, value in state_dict.items()} + torch.save(renamed, path) + + with pytest.raises(SystemExit): + upload_model.verify_checkpoint(path, spec.builder) + + +def test_verify_checkpoint_rejects_pickled_full_module(tmp_path): + spec = MODEL_REGISTRY["big_tf_unet"] + path = tmp_path / "legacy.pt" + torch.save(nn.Linear(2, 2), path) + + with pytest.raises(SystemExit): + upload_model.verify_checkpoint(path, spec.builder) From 016f152efa9f0eb54ad2c9458395dc0c6df9dab5 Mon Sep 17 00:00:00 2001 From: Nathaniel Chen Date: Sat, 4 Jul 2026 19:42:14 -0400 Subject: [PATCH 7/9] docs: quickstart-first README with CLI + auto-download; sync CLAUDE.md --- CLAUDE.md | 17 +++++--- README.md | 128 +++++++++++++++++++++++++++++++++++++----------------- 2 files changed, 100 insertions(+), 45 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 84880f6..ad3cd45 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -15,9 +15,14 @@ uv sync --dev # include dev tools (pytest, ruff, etc.) uv sync --group train # include training deps (lightning, h5py, etc.) # Run the web app -python -m TokEye.app # starts on localhost:7860 -python -m TokEye.app --port 8888 # custom port -python -m TokEye.app --share # public Gradio link +tokeye app # starts on localhost:7860 (or: python -m tokeye.app) +tokeye app --port 8888 # custom port +tokeye app --share # public Gradio link + +# Headless CLI (no gradio import; safe for HPC/CI) +tokeye run "shots/*.npy" --output-dir results # batch inference +tokeye download big_tf_unet # pre-fetch weights, print cache path +tokeye example # write a synthetic demo signal # Lint uv run ruff check . @@ -43,12 +48,12 @@ Shared building blocks in `models/modules/`: `unet.py` (base U-Net), `nn.py` (la Model I/O: input `(B, 1, H, W)` → output `(B, 2, H, W)` where channel 0 = coherent activity, channel 1 = transient activity. ### App (`app/`) -Gradio web interface launched via `python -m TokEye.app`. Three tabs: +Gradio web interface launched via `tokeye app` (console script) or `python -m tokeye.app`. Three tabs: - **Analyze** (`app/analyze/`) — load signals, compute STFT spectrograms, run inference - **Annotate** (`app/tabs/annotate.py`) — manual labeling interface - **Utilities** (`app/tabs/utilities.py`) — miscellaneous tools -Inference pipeline: `app/processing/` handles model loading, tiled inference for large spectrograms, and post-processing. +Shared core modules (used by both the app and the `tokeye` CLI) live directly under `src/tokeye/`: `hub.py` (model registry + Hugging Face auto-download), `transforms.py` (STFT), `inference.py` (model inference), `batch.py` (headless batch runner), `examples.py` (synthetic demo signal), `cli.py` (the `tokeye` console entry point). ### Training (`training/`) Multi-step data pipelines (step_0 through step_7) for preparing training data from raw signals. Two regimes: `big_tf_unet/` (original) and `big_tf_unet_multiscale/` (enhanced). Uses PyTorch Lightning. @@ -64,4 +69,4 @@ Domain-specific utilities: DIII-D tokamak helpers (`extra/D3D/`), evaluation too ## CI -GitHub Actions runs on push/PR to main: `uv run ruff check .` then `uv run pytest` across Python 3.10–3.14. +GitHub Actions runs on push/PR to main: `uv run ruff check .` then `uv run pytest` across Python 3.13–3.14. diff --git a/README.md b/README.md index 7490fa5..65ac791 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,73 @@ Expected processing time: - V100: < 0.5 seconds on any size spectrogram after warmup. - CPU: ~5-10 seconds. +## Quickstart + +```bash +pip install tokeye # or: uv tool install tokeye +tokeye app # opens web app on http://localhost:7860 +``` + +- The default model downloads automatically from Hugging Face on first use (~30 MB, cached — no manual setup). +- No data handy? Click "Load Example Signal" in the app, or generate one from the shell with `tokeye example`. + +Zero-install trial: `uvx tokeye app` runs the app without installing anything into your environment. + +## Batch processing (CLI) + +For headless / scripted use (no browser needed), run inference directly: + +```bash +tokeye run "shots/*.npy" --output-dir results +``` + +`INPUT` arguments can be files, directories (all `*.npy` files inside are used), or quoted glob patterns. Each input is interpreted by its shape: +- **1D array** — a raw time series. TokEye computes its STFT spectrogram using the flags below before running inference. +- **2D array** — a precomputed spectrogram, fed to the model directly. + +For each input file, `tokeye run` writes: +- `_mask.npy` — float32 array, shape `(2, H, W)`, sigmoid scores per pixel (channel 0 = coherent, channel 1 = transient). +- `_preview.png` — a grayscale spectrogram with the mask overlaid (green = coherent, red = transient), unless `--no-png` is passed. + +The process exit code is the number of files that failed. + +Flags: +| Flag | Default | Description | +| --- | --- | --- | +| `--model` | `big_tf_unet` | Registry name or path to a `.pt`/`.pt2` checkpoint. | +| `--output-dir` | `tokeye_output` | Directory for masks and previews. | +| `--n-fft` | `1024` | STFT window size (1D inputs only). | +| `--hop` | `256` | STFT hop size (1D inputs only). | +| `--keep-dc` | off | Keep the DC bin (dropped by default). | +| `--clip-low` / `--clip-high` | `1.0` / `99.0` | Percentile clip bounds applied to the spectrogram. | +| `--threshold` | `0.5` | Mask threshold used only for the preview PNG overlay. | +| `--no-png` | off | Skip preview PNGs; write masks only. | +| `--device` | `auto` | `cpu`, `cuda`, or `auto`. | + +The released model was trained on spectrograms built with hop=128; for closest match to the training configuration use `--hop 128`. + +On HPC clusters where compute nodes have no internet access, pre-fetch the weights on the login node, then run the batch job on the compute node: + +```bash +tokeye download big_tf_unet # on the login node; prints the cached path +tokeye run ... --model big_tf_unet # on the compute node — model is already cached +``` + +## Web app guide + +`tokeye app` (or `python -m tokeye.app`) launches a Gradio interface with three tabs: +- **Analyze** — load a signal, compute its spectrogram, run a model, and visualize the result. Guided for first-time use: the model dropdown defaults to the bundled `big_tf_unet` model, the STFT transform has working defaults, and "Load Example Signal" generates a synthetic demo signal so a brand-new user needs zero files. "Analyze" runs the whole load-model → infer → visualize pipeline in one click. View modes: Original, Enhanced (percentile-clipped amplitude), Mask (thresholded model output), Amplitude. +- **Annotate** — manually draw and save mask annotations over a read-only backdrop image. +- **Utilities** — audio-format conversion and `.npy` file inspection. + +Flags: `tokeye app [--port 7860] [--share] [--open]` — `--share` creates a public Gradio link, `--open` opens a browser tab on launch. + +If you're on a remote server (e.g. an HPC login node), forward the port over SSH instead of using `--share`: +```bash +ssh -L 7860:localhost:7860 user@remote +``` +Then open `http://localhost:7860` in your local browser. + ## Verified Datatypes - DIII-D Fast Magnetics (cite) - DIII-D CO2 Interferometer (cite) @@ -32,52 +99,28 @@ Recall Scores: With more data, comes better models. Please contribute to the project! -## Installation - -[uv](https://docs.astral.sh/uv/) (recommended) -```bash -git clone git@github.com:PlasmaControl/TokEye.git -cd TokEye -uv sync -``` - -pip (from PyPI) -```bash -pip install tokeye -``` +## Installation (from source / development) -pip (from source) +[uv](https://docs.astral.sh/uv/) is the dev tool for this repo: ```bash git clone git@github.com:PlasmaControl/TokEye.git cd TokEye -python3 -m venv .venv -source venv/bin/activate -pip install uv -uv sync +uv sync # core deps +uv sync --dev # + pytest, ruff, etc. +uv sync --group train # + training deps (lightning, h5py, etc.) ``` -## Usage -```bash -python -m TokEye.app -``` - -This will start a web app on `http://localhost:8888`. - -If you are on a remote server, you can use SSH port forwarding to access the web app on your local machine: -```bash -ssh -L 8888:localhost:7860 user@remote_server -``` -Then open your web browser and navigate to `http://localhost:8888`. +This creates a `.venv/`; activate it with `source .venv/bin/activate`, or prefix commands with `uv run`. ## Models -Pre-trained models are available at [this link](https://huggingface.co/collections/nc1/tokeye). -(soon to be deprecated) Pre-trained models are available at [this link](https://drive.google.com/drive/folders/1rXllPXB3eWhMvSIlp0CDSFx68lJOQG1u?usp=drive_link). +| Registry name | HF file | Description | +| --- | --- | --- | +| `big_tf_unet` | `big_tf_unet_251210.pt` | Transformer U-Net trained on multiscale (multiwindow, multihop) spectrograms. | -Copy them into the `models/` directory after downloading them. -- big_mode_v1.pt: Original training regime (window = 1024, hop = 128) -- big_mode_v2.pt: Trained on multiscale (multiwindow, multihop) spectrograms -- big_tf_unet_251210.pt: Trained on multiscale (multiwindow, multihop) spectrograms +Weights are hosted on [Hugging Face](https://huggingface.co/PlasmaControl/tokeye) and download automatically the first time a registry name is used (cached in `~/.cache/huggingface`). Override the source repo with the `TOKEYE_HF_REPO` environment variable. + +To use a local checkpoint instead, put `.pt`/`.pt2` files in a `model/` directory (picked up by the app's model dropdown) or pass a path directly via `--model PATH`. Input should be a tensor that has shape (B, 1, H, W) where B, H, and W can vary Output will be a tensor of shape (B, 2, H, W) @@ -88,8 +131,15 @@ The first channel of the output will return preferential measurements of coheren THe second channel of the output will return preferential measurements of transient activity ## Data -Right now, keep all data as 1d numpy float arrays. No need to normalize or preprocess them. -Copy them into the `data/` directory. +Keep signals as 1D numpy float arrays (raw time series) — no need to normalize or preprocess them. The CLI also accepts 2D arrays (precomputed spectrograms) directly. The app scans a signal directory for `.npy` files (default `data/input`, configurable in the Analyze tab). + +## Development + +```bash +uv sync --dev +uv run ruff check . +uv run pytest +``` ## Citation If you use this code in your research, please cite: @@ -105,4 +155,4 @@ If you use this code in your research, please cite: ``` ## Contact -Please check back for updates or reach out to Nathaniel Chen at nathaniel [at] princeton [dot] edu. +Nathaniel Chen — nathaniel [at] princeton [dot] edu — https://nathanielchen.net From 0edfe0b7e14d82f5f85834ace8074f220306a99d Mon Sep 17 00:00:00 2001 From: Nathaniel Chen Date: Sat, 4 Jul 2026 19:46:51 -0400 Subject: [PATCH 8/9] docs: fix typo in model output description --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 65ac791..8308480 100644 --- a/README.md +++ b/README.md @@ -128,7 +128,7 @@ Output will be a tensor of shape (B, 2, H, W) Best performance when spectrograms are oriented so that when they are plotted with matplotlib, the lowest frequency bin is oriented with the bottom when `origin='lower'`. Spectrograms should be standardized (mean = 0, std = 1). If baseline activity is very strong, clipping the input may help, but is generally not needed. The first channel of the output will return preferential measurements of coherent activity (useful for most tasks) -THe second channel of the output will return preferential measurements of transient activity +The second channel of the output will return preferential measurements of transient activity ## Data Keep signals as 1D numpy float arrays (raw time series) — no need to normalize or preprocess them. The CLI also accepts 2D arrays (precomputed spectrograms) directly. The app scans a signal directory for `.npy` files (default `data/input`, configurable in the Analyze tab). From 55e375add8e37b997ad56cecce7cf27efa5c7243 Mon Sep 17 00:00:00 2001 From: Nathaniel Chen Date: Sat, 4 Jul 2026 19:58:22 -0400 Subject: [PATCH 9/9] fix: clean CLI errors for Hugging Face hub failures; import-hygiene guard; README python note --- README.md | 1 + src/tokeye/cli.py | 24 ++++++++++++++-- tests/test_cli.py | 45 ++++++++++++++++++++++++++++++ tests/test_import_hygiene.py | 54 ++++++++++++++++++++++++++++++++++++ 4 files changed, 122 insertions(+), 2 deletions(-) create mode 100644 tests/test_import_hygiene.py diff --git a/README.md b/README.md index 8308480..1433976 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ tokeye app # opens web app on http://localhost:7860 - The default model downloads automatically from Hugging Face on first use (~30 MB, cached — no manual setup). - No data handy? Click "Load Example Signal" in the app, or generate one from the shell with `tokeye example`. +- `pip install` requires Python >= 3.13; `uvx`/`uv tool install` fetch a compatible Python automatically. Zero-install trial: `uvx tokeye app` runs the app without installing anything into your environment. diff --git a/src/tokeye/cli.py b/src/tokeye/cli.py index a490bf7..d0e5d1c 100644 --- a/src/tokeye/cli.py +++ b/src/tokeye/cli.py @@ -128,8 +128,10 @@ def _handle_app(args: argparse.Namespace) -> int: def _handle_run(args: argparse.Namespace) -> int: + from huggingface_hub.errors import HfHubHTTPError + from tokeye import batch - from tokeye.hub import DEFAULT_MODEL + from tokeye.hub import DEFAULT_MODEL, DEFAULT_REPO_ID stft_kwargs = { "n_fft": args.n_fft, @@ -153,10 +155,20 @@ def _handle_run(args: argparse.Namespace) -> int: except (ValueError, FileNotFoundError) as exc: print(f"error: {exc}", file=sys.stderr) return 2 + except (HfHubHTTPError, OSError) as exc: + print( + f"error: could not download model {model!r} from Hugging Face " + f"repo {DEFAULT_REPO_ID!r}: {exc}. If the repo has moved, set " + "TOKEYE_HF_REPO to override.", + file=sys.stderr, + ) + return 2 def _handle_download(args: argparse.Namespace) -> int: - from tokeye.hub import DEFAULT_MODEL, download_model + from huggingface_hub.errors import HfHubHTTPError + + from tokeye.hub import DEFAULT_MODEL, DEFAULT_REPO_ID, download_model names = args.models or [DEFAULT_MODEL] for name in names: @@ -165,6 +177,14 @@ def _handle_download(args: argparse.Namespace) -> int: except ValueError as exc: print(f"error: {exc}", file=sys.stderr) return 2 + except (HfHubHTTPError, OSError) as exc: + print( + f"error: could not download model {name!r} from Hugging Face " + f"repo {DEFAULT_REPO_ID!r}: {exc}. If the repo has moved, set " + "TOKEYE_HF_REPO to override.", + file=sys.stderr, + ) + return 2 print(path) return 0 diff --git a/tests/test_cli.py b/tests/test_cli.py index 2be5bc0..0f0f780 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,11 +1,25 @@ from __future__ import annotations +import httpx import numpy as np import pytest +from huggingface_hub.errors import RepositoryNotFoundError from tokeye.cli import build_parser, main +def _repository_not_found_error() -> RepositoryNotFoundError: + """Build a real RepositoryNotFoundError the way huggingface_hub does. + + ``HfHubHTTPError`` (its base class) requires a ``response`` kwarg with no + default, so a bare ``RepositoryNotFoundError("msg")`` raises ``TypeError``. + """ + response = httpx.Response( + 404, request=httpx.Request("GET", "https://huggingface.co/does/not/exist") + ) + return RepositoryNotFoundError("Repository Not Found", response=response) + + class TestBuildParser: def test_run_subcommand_parses_options(self): parser = build_parser() @@ -129,6 +143,37 @@ def test_download_unknown_model_returns_two_clean_error(self, capsys): err = capsys.readouterr().err assert "not_a_real_model_name" in err + def test_download_hub_error_returns_two_clean_error(self, monkeypatch, capsys): + def fake_download_model(name, repo_id=None): + raise _repository_not_found_error() + + monkeypatch.setattr("tokeye.hub.download_model", fake_download_model) + + exit_code = main(["download"]) + + assert exit_code == 2 + err = capsys.readouterr().err + assert "PlasmaControl/tokeye" in err + assert "TOKEYE_HF_REPO" in err + assert "Traceback" not in err + + def test_run_hub_error_returns_two_clean_error(self, tmp_path, monkeypatch, capsys): + input_path = tmp_path / "input.npy" + np.save(input_path, np.zeros((64, 32), dtype=np.float32)) + + def fake_load_model(source, device="auto"): + raise _repository_not_found_error() + + monkeypatch.setattr("tokeye.hub.load_model", fake_load_model) + + exit_code = main(["run", str(input_path), "--device", "cpu"]) + + assert exit_code == 2 + err = capsys.readouterr().err + assert "PlasmaControl/tokeye" in err + assert "TOKEYE_HF_REPO" in err + assert "Traceback" not in err + def test_app_subcommand_delegates_to_app_main(self, monkeypatch): calls = {} diff --git a/tests/test_import_hygiene.py b/tests/test_import_hygiene.py new file mode 100644 index 0000000..35e8eac --- /dev/null +++ b/tests/test_import_hygiene.py @@ -0,0 +1,54 @@ +"""HPC-safe import guarantees. + +``tokeye.batch`` must stay importable without pulling in gradio (it needs to +run on HPC login/compute nodes that may not have a display or gradio +installed). ``tokeye.cli`` must stay importable without pulling in gradio +*or* torch, so plain ``tokeye --help`` returns instantly. + +These checks run in a subprocess: importing the modules in-process (as the +rest of the test suite does throughout the run) would leave the relevant +modules in ``sys.modules`` regardless of what any single import statement +pulls in, masking a regression. +""" + +from __future__ import annotations + +import subprocess +import sys + + +def test_batch_import_does_not_pull_in_gradio(): + result = subprocess.run( + [ + sys.executable, + "-c", + "import tokeye.batch, sys; " + "assert 'gradio' not in sys.modules; " + "print('ok')", + ], + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, result.stderr + assert "ok" in result.stdout + + +def test_cli_import_does_not_pull_in_gradio_or_torch(): + result = subprocess.run( + [ + sys.executable, + "-c", + "import tokeye.cli, sys; " + "assert 'gradio' not in sys.modules; " + "assert 'torch' not in sys.modules; " + "print('ok')", + ], + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, result.stderr + assert "ok" in result.stdout