diff --git a/.gitignore b/.gitignore index d39bf03a..5db69689 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ __pycache__ /build docs/build docs/source/_autosummary +label_tool_workspace docs/source/_static/info_plot.png examples/onion_analysis/data.json examples/analysis_workflow/lens.json @@ -23,4 +24,5 @@ tests/systems/.* .pytest_cache .ruff_cache .venv +venv diff --git a/docs/source/_static/label_bar.png b/docs/source/_static/label_bar.png deleted file mode 100644 index da11d389..00000000 Binary files a/docs/source/_static/label_bar.png and /dev/null differ diff --git a/docs/source/_static/label_menu.png b/docs/source/_static/label_menu.png deleted file mode 100644 index f5b84609..00000000 Binary files a/docs/source/_static/label_menu.png and /dev/null differ diff --git a/docs/source/_static/label_tool.png b/docs/source/_static/label_tool.png index 2123ba38..7953c5dd 100644 Binary files a/docs/source/_static/label_tool.png and b/docs/source/_static/label_tool.png differ diff --git a/docs/source/label_tool/label_tool.rst b/docs/source/label_tool/label_tool.rst index 2882baec..296ad3d4 100644 --- a/docs/source/label_tool/label_tool.rst +++ b/docs/source/label_tool/label_tool.rst @@ -1,13 +1,16 @@ The Label Tool ============== -The ``dynsight label_tool`` is a simple web application that allows users to -label images. Picture labelling is a crucial step in many computer vision tasks, -such as the creation of initial training dataset to train Convolutional Neural -Networks (CNNs) model. The current version of `dynsight vision <../_autosummary/dynsight.vision.VisionInstance.html>`_ +The ``dynsight label_tool`` is a local web application for labeling images +and building training datasets. Picture labelling is a crucial step in many +computer vision tasks, such as the creation of the initial dataset used to +train Convolutional Neural Networks (CNNs). The current version of +`dynsight vision <../_autosummary/dynsight.vision.VisionInstance.html>`_ exploits the power of the `YOLO models `_ -for computer vision tasks. Thus, the ``label_tool`` has been specifically -designed to work with the YOLO dataset format. +for computer vision tasks. Thus, the ``label_tool`` writes datasets directly +in the YOLO format expected by +`set_training_dataset <../_autosummary/dynsight.vision.VisionInstance.html#dynsight.vision.VisionInstance.set_training_dataset>`_, +so they can be used for training without any manual editing. .. image:: ../_static/label_tool.png @@ -30,7 +33,10 @@ The ``label_tool`` application can be executed in 2 main ways: import dynsight - dynsight.vision.label_tool(port=8888) #port selection is optional + dynsight.vision.label_tool( + port=8888, # optional + workspace="my_workspace", # optional + ) In both cases a localhost server should start and the application should automatically appear in your default web browser. @@ -41,30 +47,53 @@ automatically appear in your default web browser. not open automatically, you can manually open it by copying and pasting the URL provided in the terminal output. +All uploaded images are stored inside the *workspace* directory +(``./label_tool_workspace`` by default). The labeling session (labels and +boxes) is kept in memory and is **never written to disk automatically**: +use the *Save session* button to write it to a JSON file at a path of your +choice, and *Load* to restore it later. If the session has unsaved changes, +the *Quit* button asks whether to save it before stopping the server +(``Ctrl+C`` in the terminal also stops it). + ------- The GUI ------- -The ``label_tool`` Graphical User Interface is divided in three main panels: +The Graphical User Interface is divided in three main panels: -* **The image panel**: where loaded images appear and labels can be drawn. +* **The labels panel** (top left): create the object classes. Each label + shows its YOLO class ID, its color and the number of boxes drawn with it. + Class IDs follow the order of this list and are stable across exports. -* **The label menu panel**: where labels can be created and edited. +* **The images panel** (bottom left): add content with ``+ Images`` or + ``+ Video`` (frames are extracted at a chosen interval), or by dragging + and dropping files onto the canvas. Each entry shows a thumbnail and its + number of annotations. -.. image:: ../_static/label_menu.png - :align: center +* **The canvas** (right): displays the current image and the bounding + boxes. -* **The commands panel**: where all the available commands can be executed. +Annotating is done directly on the canvas: -.. image:: ../_static/label_bar.png +* **Draw**: select a label, then click and drag. +* **Select**: click a box. +* **Move / resize**: drag a selected box, or drag one of its handles. +* **Change label**: select a box, then click a different label. +* **Delete**: right-click a box, or select it and press backspace. +* **Navigate**: mouse wheel to zoom, space (or middle mouse) drag to pan, + arrow keys to switch image. -Using the ``Choose File`` button, users can select the image(s) they want to -label. Once the image is loaded, users can start drawing labels by clicking and -dragging on the image panel. The label menu panel allows users to create and -edit labels. Finally, the commands panel provides a set of exporting options: +Every long operation (image and video uploads, frame extraction, dataset +export and synthesis) shows a progress bar at the bottom of the canvas. -* **Export label**: Download a single ``.txt`` file in YOLO format containing the labels for the current image. +Two export options are available in the top bar. Both write the dataset +folder directly to disk (inside the workspace by default) together with a +ready-to-use ``dataset.yaml``: -* **Export dataset**: Download a YOLO dataset from the loaded images with the labels and create the initial yaml configuration file to be used in the YOLO training process. +* **Export dataset**: exports the loaded images and their labels as a YOLO + dataset, with a configurable (and optionally shuffled) train/validation + split. -* **Synthesize dataset**: Create a synthetic dataset from the drawn labels randomizing the object position in different images (useful when a low number of images is available). +* **Synthesize**: creates a synthetic dataset by pasting the annotated + crops at random, non-overlapping positions onto uniform backgrounds + (useful when only a few labeled images are available). diff --git a/justfile b/justfile index 9269edd3..3931527a 100644 --- a/justfile +++ b/justfile @@ -1,3 +1,30 @@ +# Works with both uv and conda: +# - If a project-local uv virtualenv (./.venv) exists, its tools are +# used automatically (no activation needed). +# - Otherwise the active environment is used (conda, system, ...). +dot_venv_bin := justfile_directory() / ".venv/bin" +venv_bin := justfile_directory() / "venv/bin" +src_dir := justfile_directory() / "src" + +export PATH := if path_exists(dot_venv_bin) == "true" { + dot_venv_bin + ":" + env("PATH") +} else if path_exists(venv_bin) == "true" { + venv_bin + ":" + env("PATH") +} else { + env("PATH") +} + +# Import the package from ./src regardless of the editable install. +# This keeps the checks working even when the .pth file of the install +# is unreadable to Python, which happens on macOS when a synced folder +# (iCloud Desktop/Documents) sets the "hidden" flag on it: Python >= +# 3.11 silently skips hidden .pth files. +export PYTHONPATH := if env_var_or_default("PYTHONPATH", "") == "" { + src_dir +} else { + src_dir + ":" + env_var_or_default("PYTHONPATH", "") +} + # List all commands. default: @just --list @@ -8,33 +35,67 @@ docs: make -C docs html echo Docs are in $PWD/docs/build/html/index.html -# Do a dev install. +# Do a dev install (uv venv, conda or plain pip - autodetected). dev: - pip install -e '.[dev]' + #!/usr/bin/env bash + set -euo pipefail + # An existing project venv wins over the active environment, and + # ./.venv wins over ./venv (same order as the PATH setting above). + target="" + for candidate in .venv venv; do + if [ -d "$candidate" ]; then target="$candidate"; break; fi + done + if [ -n "$target" ] && command -v uv >/dev/null 2>&1; then + echo "Installing into ./$target with uv" + uv pip install --python "$target/bin/python" -e '.[dev]' + elif [ -n "${CONDA_PREFIX:-}" ]; then + echo "Installing into conda env '${CONDA_DEFAULT_ENV:-}' with pip" + pip install -e '.[dev]' + elif command -v uv >/dev/null 2>&1; then + echo "Creating ./.venv with uv" + uv venv + uv pip install -e '.[dev]' + else + pip install -e '.[dev]' + fi + # On macOS the .pth file of the editable install can carry the + # "hidden" flag, which makes Python >= 3.11 skip it, so that the + # package fails to import (including from the label_tool command). + # Some setups keep re-applying the flag to dot-directories such as + # ./.venv; a venv named ./venv avoids it. The recipes above do not + # depend on the .pth anyway: they import the package from ./src. + if [ "$(uname)" = "Darwin" ] && [ -n "$target" ]; then + chflags nohidden "$target"/lib/python*/site-packages/*.pth 2>/dev/null || true + fi # Run code checks. check: #!/usr/bin/env bash + # bash 3.2 (the macOS default) does not run the ERR trap when a + # subshell fails, so failures are collected explicitly: without this, + # `just check` reported success even when a step failed. error=0 - trap error=1 ERR + failed=() - echo - (set -x; ruff check . ) + run() { + echo + ( set -x; "$@" ) || { error=1; failed+=("$1"); } + } - echo - ( set -x; ruff format --check . ) + run ruff check . + run ruff format --check . + run mypy . + run pytest --cov=src --cov-report term-missing + run make -C docs doctest echo - ( set -x; mypy . ) - - echo - ( set -x; pytest --cov=src --cov-report term-missing ) - - echo - ( set -x; make -C docs doctest ) - - test $error = 0 + if [ $error -ne 0 ]; then + echo "FAILED: ${failed[*]}" + else + echo "All checks passed." + fi + exit $error # Auto-fix code issues. fix: diff --git a/pyproject.toml b/pyproject.toml index 3acfdf36..ddf2f886 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,7 +57,19 @@ line-length = 79 [tool.ruff.lint] select = ["ALL"] -ignore = ["ANN401", "COM812", "ISC001", "FBT001", "FBT002", "PLR0913", "G004"] +# CPY001: the licence lives in LICENSE, not in a header of every file. +# PLR0917: counterpart of PLR0913 for positional arguments. +ignore = [ + "ANN401", + "COM812", + "CPY001", + "ISC001", + "FBT001", + "FBT002", + "PLR0913", + "PLR0917", + "G004", +] [tool.ruff.lint.pydocstyle] convention = "google" @@ -109,7 +121,8 @@ check_untyped_defs = true disallow_untyped_decorators = true warn_unreachable = true disallow_any_generics = true -exclude = 'docs/build/html/_static' +# Directories starting with a dot (such as .venv) are skipped already. +exclude = ['docs/build/html/_static', '^venv/'] [[tool.mypy.overrides]] module = [ diff --git a/src/dynsight/_internal/analysis/spatial_average.py b/src/dynsight/_internal/analysis/spatial_average.py index 64b21870..4da97a62 100644 --- a/src/dynsight/_internal/analysis/spatial_average.py +++ b/src/dynsight/_internal/analysis/spatial_average.py @@ -132,8 +132,8 @@ def spatialaverage( Raises: ValueError: - If the input descriptor array does not have 2 or 3 dimensions, - an error is raised. + If the input descriptor array does not have 2 or 3 dimensions, or + if its number of frames does not match the (sliced) trajectory. Example: @@ -184,15 +184,24 @@ def spatialaverage( msg = "descriptor_array must have ndim == 2 or ndim == 3." raise ValueError(msg) + frame_indices = list( + range(*trajslice.indices(universe.trajectory.n_frames)) + ) + if descriptor_array.shape[1] != len(frame_indices): + msg = ( + f"descriptor_array covers {descriptor_array.shape[1]} frames, but " + f"the trajectory (after slicing) has {len(frame_indices)}. " + "Descriptors such as LENS and timeSOAP are defined on pairs of " + "frames and are one frame shorter than the trajectory they come " + "from: slice the Trj to match before averaging." + ) + raise ValueError(msg) + pool = Pool( processes=n_jobs, initializer=initworker, initargs=(shared_array, shape, dtype), ) - - frame_indices = list( - range(*trajslice.indices(universe.trajectory.n_frames)) - ) args = [ (universe, selection, r_cut, traj_frame, i, is_vector) for i, traj_frame in enumerate(frame_indices) diff --git a/src/dynsight/_internal/descriptors/misc.py b/src/dynsight/_internal/descriptors/misc.py index 66e5d644..345ff37f 100644 --- a/src/dynsight/_internal/descriptors/misc.py +++ b/src/dynsight/_internal/descriptors/misc.py @@ -12,10 +12,36 @@ from scipy.spatial.distance import cosine +def _sliced_n_frames( + universe: Universe, + trajslice: slice | None, +) -> int: + """Number of frames the given slice selects, without reading them.""" + if trajslice is None: + return len(universe.trajectory) + return len(range(*trajslice.indices(len(universe.trajectory)))) + + +def _check_neigh_list( + neigh_list_per_frame: list[list[AtomGroup]], + n_frames: int, +) -> None: + """Fail early and clearly on a neighbor list / trajectory mismatch.""" + if len(neigh_list_per_frame) != n_frames: + msg = ( + f"neigh_list_per_frame covers {len(neigh_list_per_frame)} frames, " + f"but the trajectory (after slicing) has {n_frames}. Compute the " + "neighbor list and the descriptor on the same Trj, and pass the " + "same trajslice to both." + ) + raise ValueError(msg) + + def orientational_order_param( universe: Universe, neigh_list_per_frame: list[list[AtomGroup]], order: int = 6, + trajslice: slice | None = None, ) -> NDArray[np.float64]: r"""Compute the magnitude of the orientational order parameter. @@ -39,6 +65,10 @@ def orientational_order_param( order: the order of the symmetry measured by the descriptor. Default is 6, corresponding to the hexatic order parameter. + trajslice: the slice of frames the neighbor list was computed on. Must + match the slice used to build ``neigh_list_per_frame``; if None, + the whole trajectory is used. + Returns: An array of shape (n_atoms, n_frames), with the values of psi. @@ -71,11 +101,13 @@ def orientational_order_param( """ n_atoms = universe.atoms.n_atoms - n_frames = len(universe.trajectory) + n_frames = _sliced_n_frames(universe, trajslice) + _check_neigh_list(neigh_list_per_frame, n_frames) psi = np.zeros((n_atoms, n_frames)) - for t, _ in enumerate(universe.trajectory): + frames = slice(None) if trajslice is None else trajslice + for t, _ in enumerate(universe.trajectory[frames]): frame = universe.atoms.positions[:, :2].copy() for i, atom_i in enumerate(frame): @@ -136,6 +168,7 @@ def compute_mean_alignment( def velocity_alignment( universe: Universe, neigh_list_per_frame: list[list[AtomGroup]], + trajslice: slice | None = None, ) -> NDArray[np.float64]: """Compute average velocity alignment phi. @@ -148,6 +181,10 @@ def velocity_alignment( neigh_list_per_frame: A frame-by-frame list of the neighbors of each atom, output of :func:`listNeighboursAlongTrajectory`. + trajslice: the slice of frames the neighbor list was computed on. Must + match the slice used to build ``neigh_list_per_frame``; if None, + the whole trajectory is used. + Returns: If the Universe inclused velocities, the output has shape (n_atoms, n_frames), otherwise it has shape (n_atoms, n_frames - 1). @@ -181,7 +218,9 @@ def velocity_alignment( """ n_atoms = universe.atoms.n_atoms - n_frames = len(universe.trajectory) + n_frames = _sliced_n_frames(universe, trajslice) + _check_neigh_list(neigh_list_per_frame, n_frames) + frames = slice(None) if trajslice is None else trajslice def cosine_distance( a: NDArray[np.float64], @@ -194,7 +233,7 @@ def cosine_distance( and universe.atoms.velocities is not None ): # If the Universe has velocities, use them phi = np.zeros((n_frames, n_atoms)) - for t, _ in enumerate(universe.trajectory): + for t, _ in enumerate(universe.trajectory[frames]): phi[t] = compute_mean_alignment( neigh_list_per_frame[t], vectors=universe.atoms.velocities, @@ -205,7 +244,7 @@ def cosine_distance( # If the Universe does not has velocities, use the displacements r_0 = None phi = np.zeros((n_frames - 1, n_atoms)) - for t, _ in enumerate(universe.trajectory): + for t, _ in enumerate(universe.trajectory[frames]): r_1 = universe.atoms.positions.copy() if t == 0: r_0 = r_1 diff --git a/src/dynsight/_internal/lens/lens.py b/src/dynsight/_internal/lens/lens.py index 88e083e2..079f3eff 100644 --- a/src/dynsight/_internal/lens/lens.py +++ b/src/dynsight/_internal/lens/lens.py @@ -83,7 +83,7 @@ def neighbor_list_celllist_centers( # noqa: C901, PLR0912 n_neigh = np.zeros(n_cent, dtype=np.int32) # ---- count the neighbors for each center ---- - for i in prange(n_cent): + for i in prange(n_cent): # type: ignore[attr-defined] cx = int(positions_cent[i, 0] / box[0] * nx) % nx cy = int(positions_cent[i, 1] / box[1] * ny) % ny cz = int(positions_cent[i, 2] / box[2] * nz) % nz @@ -114,7 +114,7 @@ def neighbor_list_celllist_centers( # noqa: C901, PLR0912 cursor = np.zeros(n_cent, dtype=np.int32) # ---- fill up neighbors' lists ---- - for i in prange(n_cent): + for i in prange(n_cent): # type: ignore[attr-defined] cx = int(positions_cent[i, 0] / box[0] * nx) % nx cy = int(positions_cent[i, 1] / box[1] * ny) % ny cz = int(positions_cent[i, 2] / box[2] * nz) % nz @@ -271,7 +271,10 @@ def compute_lens( coords = universe.atoms.positions mins = coords.min(axis=0) maxs = coords.max(axis=0) - box = (maxs - mins) * 1.01 + # Pad every side by r_cut: a flat system (e.g. 2D data coming + # from dynsight.vision) has zero extent along one axis, and a + # zero-length box side breaks the cell-list construction. + box = (maxs - mins) + 2 * r_cut indptr_t1, indices_t1 = neighbor_list_celllist_centers( positions_env=pos_env1, positions_cent=pos_cent1, @@ -290,7 +293,10 @@ def compute_lens( coords = universe.atoms.positions mins = coords.min(axis=0) maxs = coords.max(axis=0) - box = (maxs - mins) * 1.01 + # Pad every side by r_cut: a flat system (e.g. 2D data coming + # from dynsight.vision) has zero extent along one axis, and a + # zero-length box side breaks the cell-list construction. + box = (maxs - mins) + 2 * r_cut indptr_t2, indices_t2 = neighbor_list_celllist_centers( positions_env=pos_env2, positions_cent=pos_cent2, diff --git a/src/dynsight/_internal/track/track.py b/src/dynsight/_internal/track/track.py index 8927240d..3b902eb7 100644 --- a/src/dynsight/_internal/track/track.py +++ b/src/dynsight/_internal/track/track.py @@ -4,6 +4,7 @@ import logging from pathlib import Path +from typing import TYPE_CHECKING import pandas as pd import trackpy as tp @@ -11,21 +12,29 @@ from dynsight.trajectory import Trj from dynsight.utilities import read_xyz +if TYPE_CHECKING: + from dynsight._internal.utilities.utilities import Col + logging.basicConfig( level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s", ) logger = logging.getLogger(__name__) +# Number of columns of an atom line without the atom name. +_COORDS_ONLY = 3 +# Name given to objects read from a file without the name column. +_DEFAULT_NAME = "C" + def track_xyz( input_xyz: Path, output_xyz: Path, search_range: float, memory: int = 1, - adaptive_stop: None | float = 0.95, - adaptive_step: None | float = 0.5, -) -> Trj: + adaptive_stop: float | None = 0.95, + adaptive_step: float | None = 0.5, +) -> Trj | None: """Track particles from an ``.xyz`` file and write a new file with IDs. The input ``.xyz`` is assumed to contain only raw 3D coordinates @@ -49,8 +58,33 @@ def track_xyz( ... - The output file will have the same structure, but each line will start - with the tracked particle ID. + The output file lists the tracked particle ID at the end of each + line. Objects read from a file without the name column are written + with the placeholder name ``C``, so that the output is always a + valid ``.xyz`` file:: + + + comment line + + ... + + .. important:: + + The output file holds **one line per detection**, so its frames + contain different numbers of objects whenever the detector missed + an object or found a spurious one. Such a file records faithfully + what was detected, but it is not a trajectory: trajectory readers + require a constant number of particles. + + This function therefore returns a :class:`.trajectory.Trj` **only + when every frame contains the same number of particles**. Otherwise + it writes the file, logs how large the variation is, and returns + ``None`` rather than handing back an object that raises as soon as + a descriptor is computed on it. + + If you get ``None``, either improve the detections upstream, raise + ``memory`` so that briefly-lost objects keep their ID, or build a + trajectory yourself by keeping only the IDs present in every frame. Parameters: input_xyz: @@ -85,6 +119,11 @@ def track_xyz( Factor by which the `search_range` is multiplied to reduce it during adaptive search. Effective only if `adaptive_stop` is not `None`. + + Returns: + A :class:`.trajectory.Trj` built from the output file if every frame + holds the same number of particles, ``None`` otherwise. The output + file is written in both cases. """ if adaptive_stop is None and adaptive_step is not None: msg = "adaptive_step is set but adaptive_stop is None." @@ -101,7 +140,7 @@ def track_xyz( raise FileNotFoundError(msg) positions = read_xyz( - input_xyz=input_xyz, cols_order=["name", "x", "y", "z"] + input_xyz=input_xyz, cols_order=_detect_cols_order(input_xyz) ) if not {"frame", "x", "y", "z"}.issubset(positions.columns): @@ -130,57 +169,51 @@ def track_xyz( pid = int(row["particle"]) x, y, z = row["x"], row["y"], row["z"] name = row.get("name") - if name is not None and pd.notna(name): - f.write(f"{name} {x:.6f} {y:.6f} {z:.6f} {pid}\n") - else: - f.write(f"{x:.6f} {y:.6f} {z:.6f} {pid}\n") + if name is None or pd.isna(name): + # The name column is always written, so that the + # output stays a readable .xyz file. + name = _DEFAULT_NAME + f.write(f"{name} {x:.6f} {y:.6f} {z:.6f} {pid}\n") logger.info(f"Linked .xyz file written to: {output_xyz}") + + counts = linked.groupby("frame").size() + n_min, n_max = int(counts.min()), int(counts.max()) + if n_min != n_max: + logger.warning( + "The tracked frames hold between %d and %d objects, so the " + "output file is not a valid trajectory and no Trj is returned. " + "Improve the detections, increase 'memory', or keep only the " + "particle IDs present in every frame.", + n_min, + n_max, + ) + return None + return Trj.init_from_xyz(traj_file=output_xyz, dt=1) -def _collect_positions(input_xyz: Path) -> pd.DataFrame: - """Read the xyz file and return the positions dataset at each frame.""" - lines = input_xyz.read_text().splitlines() +def _detect_cols_order(input_xyz: Path) -> list[Col]: + """Return the column layout of the atom lines of an ``.xyz`` file. - data: list[dict[str, object]] = [] - frame = -1 - row = 0 - dimensions = 3 - for _ in range(len(lines)): - if row >= len(lines): + Both the `` `` and the `` `` layouts are + supported: the first atom line of the file decides which one is read. + """ + lines = input_xyz.read_text().splitlines() + for row, line in enumerate(lines): + # A frame starts with the number of objects and a comment line. + if not line.strip().isdigit(): + continue + if row + 2 >= len(lines): break - if lines[row].strip().isdigit(): - num_atoms = int(lines[row]) - frame += 1 - row += 2 # skip comment line. - for a in range(num_atoms): - if row + a >= len(lines): - break - parts = lines[row + a].strip().split() - if len(parts) == dimensions: - x, y, z = map(float, parts[0:3]) - data.append({"frame": frame, "x": x, "y": y, "z": z}) - elif len(parts) > dimensions: - name = parts[0] - x, y, z = map(float, parts[1:4]) - data.append( - { - "frame": frame, - "name": name, - "x": x, - "y": y, - "z": z, - } - ) - else: - msg = ( - "Invalid line format, expected 3 or 4 columns, " - f"found {len(parts)}" - ) - raise ValueError(msg) - row += num_atoms - else: - row += 1 - - return pd.DataFrame(data) + n_cols = len(lines[row + 2].split()) + if n_cols == _COORDS_ONLY: + return ["x", "y", "z"] + if n_cols > _COORDS_ONLY: + return ["name", "x", "y", "z"] + break + msg = ( + "Error in the .xyz format. Each line must be " + " or ." + ) + raise ValueError(msg) diff --git a/src/dynsight/_internal/trajectory/trajectory.py b/src/dynsight/_internal/trajectory/trajectory.py index ac121860..22a441aa 100644 --- a/src/dynsight/_internal/trajectory/trajectory.py +++ b/src/dynsight/_internal/trajectory/trajectory.py @@ -353,6 +353,7 @@ def get_orientational_op( self.universe, neigh_list_per_frame=neigcounts, order=order, + trajslice=self.trajslice, ) attr_dict = { @@ -404,6 +405,7 @@ def get_velocity_alignment( phi = dynsight.descriptors.velocity_alignment( self.universe, neigh_list_per_frame=neigcounts, + trajslice=self.trajslice, ) attr_dict = { diff --git a/src/dynsight/_internal/vision/label_tool.py b/src/dynsight/_internal/vision/label_tool.py index f3f2d62f..a451b39c 100644 --- a/src/dynsight/_internal/vision/label_tool.py +++ b/src/dynsight/_internal/vision/label_tool.py @@ -1,40 +1,747 @@ -import functools +"""Local web application for building YOLO training datasets. + +The tool starts a small HTTP server (standard library only) that serves +a single-page labeling GUI and a JSON API. Images (or video frames) are +stored inside a *workspace* directory, while the labeling session +(labels and boxes) is kept in memory and written to disk only when the +user explicitly saves it to a chosen file. Datasets are written +directly to disk in the exact layout expected by +:class:`dynsight.vision.VisionInstance`. +""" + +from __future__ import annotations + +import json import logging +import random +import re +import shutil import threading import webbrowser -from http.server import SimpleHTTPRequestHandler +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path -from socketserver import TCPServer +from typing import Any, Callable +from urllib.parse import parse_qs, urlparse + +import yaml +from PIL import Image logger = logging.getLogger(__name__) +_STATIC_DIR = Path(__file__).parent / "label_tool" +_STATIC_FILES = { + "/": ("index.html", "text/html; charset=utf-8"), + "/index.html": ("index.html", "text/html; charset=utf-8"), + "/styles.css": ("styles.css", "text/css; charset=utf-8"), + "/app.js": ("app.js", "text/javascript; charset=utf-8"), + "/logo.png": ("logo.png", "image/png"), +} + +_ProgressCallback = Callable[[int, int], None] + +_IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp"} +_VIDEO_SUFFIXES = {".mp4", ".avi", ".mov", ".mkv", ".webm"} +_UNSAFE_CHARS = re.compile(r"[^A-Za-z0-9._-]+") + +_MIN_SPLIT_IMAGES = 2 +_MAX_PLACEMENT_TRIES = 50 + + +def _safe_name(raw: str) -> str: + """Reduce a client-provided file name to a safe basename.""" + name = _UNSAFE_CHARS.sub("_", Path(raw).name) + if not re.search(r"[A-Za-z0-9]", Path(name).stem): + msg = f"Invalid file name: '{raw}'" + raise ValueError(msg) + return name + + +def _image_size(path: Path) -> tuple[int, int]: + """Return (width, height) of an image without loading pixel data.""" + with Image.open(path) as img: + return img.size + + +def _empty_session() -> dict[str, Any]: + """Return a new empty labeling session.""" + return {"labels": [], "annotations": {}} + + +def _normalize_session_path(raw: object) -> Path: + """Validate and normalize a user-provided session file path.""" + if not raw or not str(raw).strip(): + msg = "A file path is required for the session." + raise ValueError(msg) + path = Path(str(raw).strip()).expanduser() + if path.is_dir(): + path = path / "session.json" + elif path.suffix.lower() != ".json": + path = path.with_name(path.name + ".json") + return path + + +def save_session_file(session: dict[str, Any], raw_path: object) -> Path: + """Write a labeling session to an explicitly chosen file.""" + path = _normalize_session_path(raw_path) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + json.dump(session, f, indent=1) + return path + + +def load_session_file(raw_path: object) -> dict[str, Any]: + """Read a labeling session from a file.""" + path = _normalize_session_path(raw_path) + if not path.is_file(): + msg = f"Session file not found: '{path}'" + raise ValueError(msg) + with path.open("r", encoding="utf-8") as f: + data = json.load(f) + if not isinstance(data, dict): + msg = f"'{path}' is not a valid session file." + raise TypeError(msg) + return { + "labels": data.get("labels", []), + "annotations": data.get("annotations", {}), + } + + +class _Workspace: + """Filesystem-backed image storage of a labeling session.""" + + def __init__(self, root: Path) -> None: + self.root = root.resolve() + self.images_dir = self.root / "images" + self.images_dir.mkdir(parents=True, exist_ok=True) + + def list_images(self) -> list[dict[str, Any]]: + """Return metadata for every image stored in the workspace.""" + infos = [] + for path in sorted(self.images_dir.iterdir()): + if path.suffix.lower() not in _IMAGE_SUFFIXES: + continue + width, height = _image_size(path) + infos.append({"name": path.name, "width": width, "height": height}) + return infos + + def add_image(self, name: str, data: bytes) -> dict[str, Any]: + """Store an uploaded image after validating it.""" + safe = _safe_name(name) + if Path(safe).suffix.lower() not in _IMAGE_SUFFIXES: + msg = f"Unsupported image format: '{safe}'" + raise ValueError(msg) + dst = self.images_dir / safe + dst.write_bytes(data) + try: + width, height = _image_size(dst) + # Broad catch: PIL.Image.open may be monkey-patched by other + # libraries (e.g. ultralytics) and raise unexpected errors. + except Exception: # noqa: BLE001 + dst.unlink(missing_ok=True) + msg = f"'{safe}' is not a readable image." + raise ValueError(msg) from None + return {"name": safe, "width": width, "height": height} + + def add_video( + self, + name: str, + data: bytes, + stride: int, + on_progress: _ProgressCallback | None = None, + ) -> list[dict[str, Any]]: + """Extract frames from an uploaded video into the workspace.""" + import cv2 # noqa: PLC0415 (heavy import, only needed here) + + safe = _safe_name(name) + if Path(safe).suffix.lower() not in _VIDEO_SUFFIXES: + msg = f"Unsupported video format: '{safe}'" + raise ValueError(msg) + stride = max(1, stride) + tmp = self.root / f"_upload_{safe}" + tmp.write_bytes(data) + stem = Path(safe).stem + frames: list[dict[str, Any]] = [] + try: + capture = cv2.VideoCapture(str(tmp)) + if not capture.isOpened(): + msg = f"Could not open video '{safe}'." + raise ValueError(msg) + total = max(0, int(capture.get(cv2.CAP_PROP_FRAME_COUNT))) + index = 0 + while True: + ok, frame = capture.read() + if not ok: + break + if index % stride == 0: + frame_name = f"{stem}_{index:06d}.jpg" + cv2.imwrite(str(self.images_dir / frame_name), frame) + height, width = frame.shape[:2] + frames.append( + { + "name": frame_name, + "width": int(width), + "height": int(height), + } + ) + index += 1 + if on_progress is not None: + on_progress(index, total) + capture.release() + finally: + tmp.unlink(missing_ok=True) + if not frames: + msg = f"No frames could be extracted from '{safe}'." + raise ValueError(msg) + return frames + + def delete_image(self, name: str) -> None: + """Remove an image from the workspace.""" + (self.images_dir / _safe_name(name)).unlink(missing_ok=True) + + +def _yolo_lines( + boxes: list[dict[str, Any]], + class_ids: dict[str, int], + width: int, + height: int, +) -> str: + """Convert pixel-space boxes to YOLO txt content.""" + lines = [] + for box in boxes: + label = box["label"] + if label not in class_ids: + continue + cx = (box["x"] + box["w"] / 2) / width + cy = (box["y"] + box["h"] / 2) / height + w = box["w"] / width + h = box["h"] / height + cx, cy = min(max(cx, 0.0), 1.0), min(max(cy, 0.0), 1.0) + w, h = min(max(w, 0.0), 1.0), min(max(h, 0.0), 1.0) + lines.append(f"{class_ids[label]} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}") + return "".join(f"{line}\n" for line in lines) + + +def _dataset_dirs(dataset_path: Path) -> dict[str, Path]: + """Create and return the YOLO dataset directory layout.""" + dirs = { + "images/train": dataset_path / "images" / "train", + "images/val": dataset_path / "images" / "val", + "labels/train": dataset_path / "labels" / "train", + "labels/val": dataset_path / "labels" / "val", + } + for path in dirs.values(): + path.mkdir(parents=True, exist_ok=True) + return dirs + + +def _write_dataset_yaml(dataset_path: Path, names: list[str]) -> Path: + """Write the dataset.yaml file consumed by ``set_training_dataset``.""" + yaml_path = dataset_path / "dataset.yaml" + content = { + "path": str(dataset_path.resolve()), + "train": "images/train", + "val": "images/val", + "nc": len(names), + "names": names, + } + with yaml_path.open("w", encoding="utf-8") as f: + yaml.safe_dump(content, f, sort_keys=False) + return yaml_path + + +def _split_count(total: int, train_split: float) -> int: + """Number of training items for a given split fraction.""" + num_train = round(total * train_split) + if total >= _MIN_SPLIT_IMAGES: + num_train = min(max(num_train, 1), total - 1) + return num_train + + +def export_dataset( + workspace: _Workspace, + session: dict[str, Any], + name: str, + train_split: float = 0.8, + shuffle: bool = True, + seed: int | None = None, + output_dir: Path | None = None, + on_progress: _ProgressCallback | None = None, +) -> dict[str, Any]: + """Write a YOLO dataset from the current session to disk. + + Class IDs follow the order of the session label list, so they are + stable across exports. Every image receives a label file (empty if + it has no annotations) and every known class appears in + ``dataset.yaml`` even when unused. + """ + if not 0.0 < train_split < 1.0: + msg = "train_split must be between 0 and 1." + raise ValueError(msg) + images = workspace.list_images() + if not images: + msg = "No images in the workspace." + raise ValueError(msg) + names = [label["name"] for label in session.get("labels", [])] + if not names: + msg = "No labels defined." + raise ValueError(msg) + class_ids = {label: idx for idx, label in enumerate(names)} + annotations: dict[str, Any] = session.get("annotations", {}) + + base = output_dir if output_dir is not None else workspace.root + dataset_path = (base / _safe_name(name)).resolve() + dirs = _dataset_dirs(dataset_path) + + if shuffle: + random.Random(seed).shuffle(images) # noqa: S311 + num_train = _split_count(len(images), train_split) + + for idx, info in enumerate(images): + subset = "train" if idx < num_train else "val" + src = workspace.images_dir / info["name"] + shutil.copy2(src, dirs[f"images/{subset}"] / info["name"]) + txt = _yolo_lines( + annotations.get(info["name"], []), + class_ids, + info["width"], + info["height"], + ) + lbl = dirs[f"labels/{subset}"] / (Path(info["name"]).stem + ".txt") + lbl.write_text(txt, encoding="utf-8") + if on_progress is not None: + on_progress(idx + 1, len(images)) + + yaml_path = _write_dataset_yaml(dataset_path, names) + return { + "path": str(dataset_path), + "yaml": str(yaml_path), + "num_train": num_train, + "num_val": len(images) - num_train, + } + + +def _place_crop( + rng: random.Random, + crop_size: tuple[int, int], + canvas_size: tuple[int, int], + placed: list[tuple[float, float, float, float]], + scale_range: tuple[float, float], +) -> tuple[int, int, int, int] | None: + """Find a non-overlapping position for a crop, or ``None``.""" + for _ in range(_MAX_PLACEMENT_TRIES): + scale = rng.uniform(*scale_range) + w = max(1, int(crop_size[0] * scale)) + h = max(1, int(crop_size[1] * scale)) + if w >= canvas_size[0] or h >= canvas_size[1]: + continue + x = rng.randint(0, canvas_size[0] - w) + y = rng.randint(0, canvas_size[1] - h) + overlap = any( + x < px + pw and x + w > px and y < py + ph and y + h > py + for px, py, pw, ph in placed + ) + if not overlap: + return x, y, w, h + return None + + +def synthesize_dataset( + workspace: _Workspace, + session: dict[str, Any], + name: str, + num_images: int = 10, + width: int = 640, + height: int = 640, + per_image: int = 10, + train_split: float = 0.8, + scale_range: tuple[float, float] = (1.0, 1.0), + background: str = "#ffffff", + seed: int | None = None, + output_dir: Path | None = None, + on_progress: _ProgressCallback | None = None, +) -> dict[str, Any]: + """Generate a synthetic YOLO dataset from the annotated crops. + + Annotated regions are cut out of the source images and pasted at + random non-overlapping positions onto uniform background canvases. + """ + names = [label["name"] for label in session.get("labels", [])] + class_ids = {label: idx for idx, label in enumerate(names)} + annotations: dict[str, Any] = session.get("annotations", {}) + crops = [ + {"image": image_name, **box} + for image_name, boxes in annotations.items() + for box in boxes + if box["label"] in class_ids + and (workspace.images_dir / image_name).is_file() + ] + if not crops: + msg = "No annotations available to synthesize from." + raise ValueError(msg) + + base = output_dir if output_dir is not None else workspace.root + dataset_path = (base / _safe_name(name)).resolve() + dirs = _dataset_dirs(dataset_path) + rng = random.Random(seed) # noqa: S311 + num_train = _split_count(num_images, train_split) + + sources: dict[str, Image.Image] = {} + for idx in range(num_images): + canvas = Image.new("RGB", (width, height), background) + placed: list[tuple[float, float, float, float]] = [] + boxes: list[dict[str, Any]] = [] + for _ in range(per_image): + crop = rng.choice(crops) + if crop["image"] not in sources: + src_path = workspace.images_dir / crop["image"] + sources[crop["image"]] = Image.open(src_path).convert("RGB") + source = sources[crop["image"]] + left, top = int(crop["x"]), int(crop["y"]) + cw = max(1, int(crop["w"])) + ch = max(1, int(crop["h"])) + patch = source.crop((left, top, left + cw, top + ch)) + spot = _place_crop( + rng, (cw, ch), (width, height), placed, scale_range + ) + if spot is None: + continue + x, y, w, h = spot + canvas.paste(patch.resize((w, h)), (x, y)) + placed.append((x, y, w, h)) + boxes.append( + {"label": crop["label"], "x": x, "y": y, "w": w, "h": h} + ) + subset = "train" if idx < num_train else "val" + canvas.save(dirs[f"images/{subset}"] / f"synt_{idx:05d}.jpg") + txt = _yolo_lines(boxes, class_ids, width, height) + lbl = dirs[f"labels/{subset}"] / f"synt_{idx:05d}.txt" + lbl.write_text(txt, encoding="utf-8") + if on_progress is not None: + on_progress(idx + 1, num_images) + + for source in sources.values(): + source.close() + yaml_path = _write_dataset_yaml(dataset_path, names) + return { + "path": str(dataset_path), + "yaml": str(yaml_path), + "num_train": num_train, + "num_val": num_images - num_train, + } + + +class _LabelToolServer(ThreadingHTTPServer): + """HTTP server carrying the state shared by all requests. + + The labeling session (labels and boxes) lives in ``self.session`` + and is written to disk only through the explicit save endpoint. + ``self.progress`` mirrors the state of the long-running operation + currently in flight (if any) and is polled by the GUI to render + progress bars. + """ + + allow_reuse_address = True + daemon_threads = True + + def __init__(self, port: int, workspace: _Workspace) -> None: + self.workspace = workspace + self.session = _empty_session() + self.session_path: str | None = None + self.session_dirty = False + self.progress: dict[str, Any] = {"active": False} + self.progress_lock = threading.Lock() + super().__init__(("127.0.0.1", port), _RequestHandler) + + def start_progress(self, label: str) -> _ProgressCallback: + """Mark a long operation as active and return its callback.""" + with self.progress_lock: + self.progress = { + "active": True, + "label": label, + "done": 0, + "total": 0, + } + + def on_progress(done: int, total: int) -> None: + with self.progress_lock: + if self.progress.get("active"): + self.progress["done"] = done + self.progress["total"] = total + + return on_progress + + def end_progress(self) -> None: + """Mark the current long operation as finished.""" + with self.progress_lock: + self.progress = {"active": False} + + +class _RequestHandler(BaseHTTPRequestHandler): + """Routes static files and the JSON API.""" + + server: _LabelToolServer -class HTTPRequestHandler(SimpleHTTPRequestHandler): def log_message(self, fmt: str, *args: object) -> None: - pass + """Silence default request logging.""" + + @property + def _workspace(self) -> _Workspace: + return self.server.workspace + + def _send_json( + self, payload: dict[str, Any], status: int = HTTPStatus.OK + ) -> None: + body = json.dumps(payload).encode("utf-8") + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def _send_file( + self, path: Path, content_type: str, cache_control: str = "no-store" + ) -> None: + if not path.is_file(): + self.send_error(HTTPStatus.NOT_FOUND) + return + body = path.read_bytes() + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(body))) + self.send_header("Cache-Control", cache_control) + self.end_headers() + self.wfile.write(body) + + def _read_body(self) -> bytes: + length = int(self.headers.get("Content-Length", "0")) + return self.rfile.read(length) + + def _query(self) -> dict[str, str]: + parsed = parse_qs(urlparse(self.path).query) + return {key: values[0] for key, values in parsed.items()} + + def do_GET(self) -> None: + """Serve the GUI, workspace images and the state endpoint.""" + route = urlparse(self.path).path + if route in _STATIC_FILES: + file_name, content_type = _STATIC_FILES[route] + self._send_file(_STATIC_DIR / file_name, content_type) + elif route.startswith("/images/"): + name = _safe_name(route[len("/images/") :]) + suffix = Path(name).suffix.lower().lstrip(".") + content_type = f"image/{'jpeg' if suffix == 'jpg' else suffix}" + self._send_file( + self._workspace.images_dir / name, + content_type, + cache_control="max-age=300", + ) + elif route == "/api/state": + self._api(self._handle_state) + elif route == "/api/progress": + self._api(self._handle_progress) + else: + self.send_error(HTTPStatus.NOT_FOUND) - # do_POST must be uppercase def do_POST(self) -> None: - if self.path == "/shutdown": - self.send_response(200) - self.end_headers() - logger.info("Shutdown request received.") - threading.Thread(target=self.server.shutdown).start() + """Dispatch API mutations.""" + route = urlparse(self.path).path + handlers = { + "/api/sync": self._handle_sync, + "/api/session": self._handle_save_session, + "/api/session/load": self._handle_load_session, + "/api/images": self._handle_upload_image, + "/api/video": self._handle_upload_video, + "/api/export": self._handle_export, + "/api/synthesize": self._handle_synthesize, + "/api/shutdown": self._handle_shutdown, + } + handler = handlers.get(route) + if handler is None: + self.send_error(HTTPStatus.NOT_FOUND) + return + self._api(handler) + + def do_DELETE(self) -> None: + """Delete a workspace image.""" + route = urlparse(self.path).path + if route == "/api/images": + self._api(self._handle_delete_image) else: - self.send_error(404) + self.send_error(HTTPStatus.NOT_FOUND) + def _api(self, handler: Any) -> None: + try: + payload = handler() + except (ValueError, KeyError, TypeError, OSError) as e: + logger.warning(f"Request failed: {e}") + self._send_json({"error": str(e)}, status=HTTPStatus.BAD_REQUEST) + else: + self._send_json(payload) -class ReusableTCPServer(TCPServer): - allow_reuse_address = True + def _json_body(self) -> dict[str, Any]: + data: dict[str, Any] = json.loads(self._read_body() or b"{}") + return data + + def _session_from(self, body: dict[str, Any]) -> dict[str, Any]: + return { + "labels": body.get("labels", []), + "annotations": body.get("annotations", {}), + } + + def _handle_state(self) -> dict[str, Any]: + session = self.server.session + return { + "workspace": str(self._workspace.root), + "images": self._workspace.list_images(), + "labels": session.get("labels", []), + "annotations": session.get("annotations", {}), + "session_path": self.server.session_path, + "dirty": self.server.session_dirty, + } + + def _handle_progress(self) -> dict[str, Any]: + with self.server.progress_lock: + return dict(self.server.progress) + + def _handle_sync(self) -> dict[str, Any]: + """Update the in-memory session (no disk write).""" + self.server.session = self._session_from(self._json_body()) + self.server.session_dirty = True + return {"synced": True} + + def _handle_save_session(self) -> dict[str, Any]: + """Write the session to an explicitly chosen file.""" + body = self._json_body() + if "labels" in body or "annotations" in body: + self.server.session = self._session_from(body) + path = save_session_file(self.server.session, body.get("path")) + self.server.session_path = str(path) + self.server.session_dirty = False + return {"path": str(path)} + + def _handle_load_session(self) -> dict[str, Any]: + """Load a session file into memory and return it.""" + body = self._json_body() + session = load_session_file(body.get("path")) + self.server.session = session + self.server.session_path = str( + _normalize_session_path(body.get("path")) + ) + self.server.session_dirty = False + return session + + def _handle_upload_image(self) -> dict[str, Any]: + query = self._query() + return self._workspace.add_image(query["name"], self._read_body()) + + def _handle_upload_video(self) -> dict[str, Any]: + query = self._query() + data = self._read_body() + on_progress = self.server.start_progress("Extracting frames") + try: + frames = self._workspace.add_video( + query["name"], + data, + stride=int(query.get("stride", "1")), + on_progress=on_progress, + ) + finally: + self.server.end_progress() + return {"frames": frames} + + def _handle_delete_image(self) -> dict[str, Any]: + query = self._query() + self._workspace.delete_image(query["name"]) + return {"deleted": True} + + def _handle_export(self) -> dict[str, Any]: + body = self._json_body() + output = body.get("output_dir") + on_progress = self.server.start_progress("Exporting dataset") + try: + return export_dataset( + self._workspace, + self.server.session, + name=body.get("name", "yolo_dataset"), + train_split=float(body.get("train_split", 0.8)), + shuffle=bool(body.get("shuffle", True)), + seed=body.get("seed"), + output_dir=Path(output) if output else None, + on_progress=on_progress, + ) + finally: + self.server.end_progress() + + def _handle_synthesize(self) -> dict[str, Any]: + body = self._json_body() + output = body.get("output_dir") + on_progress = self.server.start_progress("Synthesizing dataset") + try: + return synthesize_dataset( + self._workspace, + self.server.session, + name=body.get("name", "synt_dataset"), + num_images=int(body.get("num_images", 10)), + width=int(body.get("width", 640)), + height=int(body.get("height", 640)), + per_image=int(body.get("per_image", 10)), + train_split=float(body.get("train_split", 0.8)), + scale_range=( + float(body.get("scale_min", 1.0)), + float(body.get("scale_max", 1.0)), + ), + background=str(body.get("background", "#ffffff")), + seed=body.get("seed"), + output_dir=Path(output) if output else None, + on_progress=on_progress, + ) + finally: + self.server.end_progress() + + def _handle_shutdown(self) -> dict[str, Any]: + logger.info("Shutdown requested from the GUI.") + threading.Thread(target=self.server.shutdown).start() + return {"shutdown": True} + + +def label_tool( + port: int = 8888, + workspace: str | Path | None = None, + open_browser: bool = True, +) -> None: + """Start the dynsight labeling tool. + + The tool opens in the default web browser. Uploaded images are + stored in ``workspace``, while the labeling session (labels and + boxes) is saved to disk only when explicitly requested from the + GUI, to a file path chosen by the user. The server stops with the + *Quit* button in the GUI or with ``Ctrl+C`` in the terminal. + + Parameters: + port: + Port for the local HTTP server. + workspace: + Directory where images and exported datasets are stored by + default. Defaults to ``./label_tool_workspace``. -def label_tool(port: int = 8888) -> None: - web_dir = Path(__file__).parent / "label_tool" - handler = functools.partial(HTTPRequestHandler, directory=str(web_dir)) - with ReusableTCPServer(("", port), handler) as httpd: - url = f"http://localhost:{port}/index.html" - logger.info(f"Starting server at {url}") + open_browser: + Automatically open the GUI in the default browser. + """ + root = ( + Path(workspace) if workspace else Path.cwd() / ("label_tool_workspace") + ) + server = _LabelToolServer(port, _Workspace(root)) + url = f"http://127.0.0.1:{port}/" + logger.info(f"Labeling tool running at {url}") + logger.info(f"Workspace: {root.resolve()}") + if open_browser: webbrowser.open(url) - httpd.serve_forever() - httpd.server_close() + try: + server.serve_forever() + except KeyboardInterrupt: + logger.info("Interrupted.") + finally: + server.server_close() logger.info("Server closed.") diff --git a/src/dynsight/_internal/vision/label_tool/app.js b/src/dynsight/_internal/vision/label_tool/app.js new file mode 100644 index 00000000..a550b02f --- /dev/null +++ b/src/dynsight/_internal/vision/label_tool/app.js @@ -0,0 +1,1259 @@ +"use strict"; + +/* ========================================================= + * dynsight label tool - frontend + * + * All annotation coordinates are stored in natural image + * pixels: {label, x, y, w, h} with (x, y) = top-left corner. + * The session lives in memory (mirrored to the server) and is + * written to disk only via the explicit Save session dialog. + * ========================================================= */ + +/* ---------- constants ---------- */ + +const PALETTE = [ + "#f43f5e", "#f97316", "#eab308", "#22c55e", "#06b6d4", + "#3b82f6", "#8b5cf6", "#ec4899", "#14b8a6", "#a3e635", +]; +const MIN_BOX_SIZE = 3; // px, in image space +const HANDLE_SIZE = 7; // px, in screen space +const HANDLE_HIT = 6; // px tolerance +const MIN_SCALE = 0.05; +const MAX_SCALE = 32; + +/* ---------- state ---------- */ + +const state = { + workspace: "", + images: [], // [{name, width, height}] + annotations: {}, // name -> [{label, x, y, w, h}] + labels: [], // [{name, color}] + activeLabel: null, + current: -1, + selection: -1, + view: { scale: 1, x: 0, y: 0 }, + fitted: true, // refit on container resize until the user zooms/pans +}; + +let drag = null; // {mode, ...} while a pointer drag is active +let hover = { box: -1, handle: -1 }; +let pointer = { x: 0, y: 0, inside: false }; +let spaceDown = false; +let syncTimer = null; +let dirty = false; // changes not yet saved to a session file +let sessionPath = null; // last file the session was saved to / loaded from +let quitAfterSave = false; + +const imageCache = new Map(); // name -> HTMLImageElement +const imageVersion = new Map(); // name -> int, bumped on re-upload + +function imageUrl(name) { + const version = imageVersion.get(name); + const suffix = version ? `?v=${version}` : ""; + return `/images/${encodeURIComponent(name)}${suffix}`; +} + +/* ---------- dom ---------- */ + +const $ = (id) => document.getElementById(id); +const canvas = $("canvas"); +const ctx = canvas.getContext("2d"); +const stage = $("stage"); + +/* ---------- helpers ---------- */ + +function currentImage() { + return state.images[state.current] || null; +} + +function currentBoxes() { + const img = currentImage(); + if (!img) return []; + if (!state.annotations[img.name]) state.annotations[img.name] = []; + return state.annotations[img.name]; +} + +function labelColor(name) { + const label = state.labels.find((l) => l.name === name); + return label ? label.color : "#9ca3af"; +} + +function clamp(value, lo, hi) { + return Math.max(lo, Math.min(hi, value)); +} + +function screenToImage(sx, sy) { + return { + x: (sx - state.view.x) / state.view.scale, + y: (sy - state.view.y) / state.view.scale, + }; +} + +/* ---------- api ---------- */ + +async function api(path, options = {}) { + const response = await fetch(path, options); + let payload = {}; + try { + payload = await response.json(); + } catch { + /* non-json error */ + } + if (!response.ok) { + throw new Error(payload.error || `Request failed (${response.status})`); + } + return payload; +} + +function setSaveStatus() { + const el = $("saveStatus"); + if (dirty) { + el.textContent = "● Unsaved session"; + el.className = "busy"; + } else if (sessionPath) { + el.textContent = `Saved ✓ (${sessionPath})`; + el.className = "ok"; + } else { + el.textContent = ""; + el.className = ""; + } +} + +// The session is never written to disk automatically: edits are only +// mirrored to the server's memory so a page reload does not lose work +// while the server is running. Disk writes happen exclusively through +// the "Save session" dialog, to a user-chosen path. +function markChanged() { + dirty = true; + setSaveStatus(); + clearTimeout(syncTimer); + syncTimer = setTimeout(syncSession, 300); +} + +async function syncSession() { + clearTimeout(syncTimer); + try { + await api("/api/sync", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: sessionBody(), + }); + } catch { + /* retried on the next change */ + } +} + +function sessionBody() { + return JSON.stringify({ + labels: state.labels, + annotations: state.annotations, + }); +} + +// Mirror unsaved work to the server's memory when the page is closed +// or reloaded (no disk write). The server keeps running: it is stopped +// only via the Quit button or Ctrl+C. +window.addEventListener("pagehide", () => { + navigator.sendBeacon( + "/api/sync", + new Blob([sessionBody()], { type: "application/json" }), + ); +}); + +// Warn before leaving the page with an unsaved session. +window.addEventListener("beforeunload", (e) => { + if (dirty) e.preventDefault(); +}); + +/* ---------- toasts ---------- */ + +function toast(message, cls = "", detail = "", timeout = 6000) { + const el = document.createElement("div"); + el.className = `toast ${cls}`; + el.textContent = message; + if (detail) { + const line = document.createElement("span"); + line.className = "mono"; + line.textContent = detail; + el.appendChild(line); + } + el.onclick = () => el.remove(); + $("toasts").appendChild(el); + setTimeout(() => el.remove(), timeout); +} + +/* ---------- progress ---------- */ + +let progressPoll = null; + +function showProgress(label) { + $("progressLabel").textContent = label; + $("progressPct").textContent = ""; + $("progressFill").classList.add("indeterminate"); + $("progress").classList.remove("hidden"); +} + +function setProgress(done, total) { + const fill = $("progressFill"); + if (total > 0) { + const pct = Math.min(100, Math.round((done / total) * 100)); + fill.classList.remove("indeterminate"); + fill.style.width = `${pct}%`; + $("progressPct").textContent = `${pct}%`; + } else { + fill.classList.add("indeterminate"); + $("progressPct").textContent = done > 0 ? String(done) : ""; + } +} + +function hideProgress() { + stopProgressPoll(); + $("progress").classList.add("hidden"); + $("progressFill").style.width = "0%"; + $("progressFill").classList.remove("indeterminate"); +} + +// Long server-side operations (export, synthesize, frame extraction) +// report their progress through /api/progress, polled while the main +// request is in flight. +function startProgressPoll(fallbackLabel) { + stopProgressPoll(); + progressPoll = setInterval(async () => { + try { + const p = await api("/api/progress"); + if (p.active) { + $("progressLabel").textContent = + (p.label || fallbackLabel) + "…"; + setProgress(p.done || 0, p.total || 0); + } + } catch { + /* server busy or gone; keep the bar as-is */ + } + }, 250); +} + +function stopProgressPoll() { + if (progressPoll) { + clearInterval(progressPoll); + progressPoll = null; + } +} + +/* ---------- sidebar: labels ---------- */ + +function renderLabels() { + const list = $("labelList"); + list.innerHTML = ""; + state.labels.forEach((label, idx) => { + const li = document.createElement("li"); + if (label.name === state.activeLabel) li.classList.add("active"); + + const id = document.createElement("span"); + id.className = "class-id"; + id.textContent = String(idx); + + const dot = document.createElement("span"); + dot.className = "color-dot"; + dot.style.backgroundColor = label.color; + + const name = document.createElement("span"); + name.className = "item-name"; + name.textContent = label.name; + + const count = document.createElement("span"); + count.className = "item-badge"; + count.textContent = String(countBoxes(label.name)); + + const del = document.createElement("button"); + del.className = "del-btn"; + del.textContent = "×"; + del.title = "Delete label and its boxes"; + del.onclick = (e) => { + e.stopPropagation(); + deleteLabel(label.name); + }; + + li.append(id, dot, name, count, del); + li.onclick = () => { + state.activeLabel = label.name; + const boxes = currentBoxes(); + if (state.selection >= 0 && boxes[state.selection]) { + boxes[state.selection].label = label.name; + markChanged(); + } + renderLabels(); + render(); + }; + list.appendChild(li); + }); + $("labelHint").classList.toggle("hidden", state.labels.length > 0); +} + +function countBoxes(labelName) { + let total = 0; + for (const boxes of Object.values(state.annotations)) { + total += boxes.filter((b) => b.label === labelName).length; + } + return total; +} + +function addLabel(name) { + if (!name || state.labels.some((l) => l.name === name)) return; + const color = PALETTE[state.labels.length % PALETTE.length]; + state.labels.push({ name, color }); + state.activeLabel = name; + markChanged(); + renderLabels(); + render(); +} + +function deleteLabel(name) { + const used = countBoxes(name); + if ( + used > 0 && + !confirm(`Delete label "${name}" and its ${used} box(es)?`) + ) { + return; + } + state.labels = state.labels.filter((l) => l.name !== name); + for (const key of Object.keys(state.annotations)) { + state.annotations[key] = state.annotations[key].filter( + (b) => b.label !== name, + ); + } + if (state.activeLabel === name) state.activeLabel = null; + state.selection = -1; + markChanged(); + renderLabels(); + renderImages(); + render(); +} + +$("labelForm").onsubmit = (e) => { + e.preventDefault(); + addLabel($("labelInput").value.trim()); + $("labelInput").value = ""; +}; + +/* ---------- sidebar: images ---------- */ + +function renderImages() { + const list = $("imageList"); + list.innerHTML = ""; + state.images.forEach((info, idx) => { + const li = document.createElement("li"); + if (idx === state.current) li.classList.add("active"); + + const thumb = document.createElement("img"); + thumb.className = "thumb"; + thumb.loading = "lazy"; + thumb.src = imageUrl(info.name); + + const name = document.createElement("span"); + name.className = "item-name"; + name.textContent = info.name; + name.title = info.name; + + const count = document.createElement("span"); + count.className = "item-badge"; + const n = (state.annotations[info.name] || []).length; + count.textContent = n > 0 ? String(n) : ""; + + const del = document.createElement("button"); + del.className = "del-btn"; + del.textContent = "×"; + del.title = "Remove image"; + del.onclick = (e) => { + e.stopPropagation(); + deleteImage(info.name); + }; + + li.append(thumb, name, count, del); + li.onclick = () => selectImage(idx); + list.appendChild(li); + }); + $("imageCounter").textContent = state.images.length + ? `${state.current + 1} / ${state.images.length}` + : "0 / 0"; + $("emptyState").classList.toggle("hidden", state.images.length > 0); +} + +async function deleteImage(name) { + if (!confirm(`Remove "${name}" and its annotations?`)) return; + try { + await api(`/api/images?name=${encodeURIComponent(name)}`, { + method: "DELETE", + }); + } catch (err) { + toast(`Could not delete: ${err.message}`, "error"); + return; + } + const idx = state.images.findIndex((i) => i.name === name); + state.images = state.images.filter((i) => i.name !== name); + delete state.annotations[name]; + imageCache.delete(name); + if (state.current >= state.images.length) { + state.current = state.images.length - 1; + } else if (idx <= state.current) { + state.current = Math.max(0, state.current - (idx < state.current)); + } + state.selection = -1; + markChanged(); + selectImage(state.current, true); + renderLabels(); +} + +function selectImage(idx, force = false) { + if (idx === state.current && !force) return; + state.current = clamp(idx, -1, state.images.length - 1); + state.selection = -1; + const info = currentImage(); + if (info && !imageCache.has(info.name)) { + const img = new Image(); + img.onload = () => { + if (currentImage() === info) fitView(); + }; + img.src = imageUrl(info.name); + imageCache.set(info.name, img); + } + fitView(); + renderImages(); +} + +$("prevBtn").onclick = () => { + if (state.current > 0) selectImage(state.current - 1); +}; +$("nextBtn").onclick = () => { + if (state.current < state.images.length - 1) { + selectImage(state.current + 1); + } +}; + +/* ---------- uploads ---------- */ + +async function uploadImages(files) { + const list = Array.from(files); + if (!list.length) return; + let done = 0; + showProgress(`Uploading images (0/${list.length})…`); + setProgress(0, list.length); + try { + for (const [idx, file] of list.entries()) { + try { + const info = await api( + `/api/images?name=${encodeURIComponent(file.name)}`, + { method: "POST", body: file }, + ); + const existing = state.images.findIndex( + (i) => i.name === info.name, + ); + if (existing >= 0) { + state.images[existing] = info; + imageCache.delete(info.name); + imageVersion.set( + info.name, + (imageVersion.get(info.name) || 0) + 1, + ); + } else { + state.images.push(info); + } + done += 1; + } catch (err) { + toast(`"${file.name}": ${err.message}`, "error"); + } + $("progressLabel").textContent = + `Uploading images (${idx + 1}/${list.length})…`; + setProgress(idx + 1, list.length); + } + } finally { + hideProgress(); + } + if (done > 0) { + toast(`Added ${done} image(s).`, "ok"); + if (state.current < 0) selectImage(0); + renderImages(); + } +} + +let pendingVideo = null; + +function askVideoStride(file) { + pendingVideo = file; + $("videoFileName").textContent = file.name; + $("videoDialog").showModal(); +} + +// Upload with XMLHttpRequest to get byte-level upload progress, then +// poll /api/progress while the server extracts frames. +function uploadVideo(url, file) { + return new Promise((resolve, reject) => { + const xhr = new XMLHttpRequest(); + xhr.open("POST", url); + xhr.upload.onprogress = (e) => { + if (e.lengthComputable) { + $("progressLabel").textContent = "Uploading video…"; + setProgress(e.loaded, e.total); + } + }; + xhr.upload.onload = () => { + showProgress("Extracting frames…"); + startProgressPoll("Extracting frames"); + }; + xhr.onload = () => { + let payload = {}; + try { + payload = JSON.parse(xhr.responseText); + } catch { + /* non-json */ + } + if (xhr.status >= 200 && xhr.status < 300) resolve(payload); + else { + reject( + new Error( + payload.error || `Request failed (${xhr.status})`, + ), + ); + } + }; + xhr.onerror = () => reject(new Error("Network error")); + xhr.send(file); + }); +} + +$("videoForm").onsubmit = async (e) => { + e.preventDefault(); + const stride = $("videoForm").elements.stride.value || "1"; + const file = pendingVideo; + pendingVideo = null; + $("videoDialog").close(); + if (!file) return; + showProgress("Uploading video…"); + try { + const result = await uploadVideo( + `/api/video?name=${encodeURIComponent(file.name)}` + + `&stride=${encodeURIComponent(stride)}`, + file, + ); + for (const info of result.frames) { + if (!state.images.some((i) => i.name === info.name)) { + state.images.push(info); + } + } + toast(`Added ${result.frames.length} frame(s).`, "ok"); + if (state.current < 0) selectImage(0); + renderImages(); + } catch (err) { + toast(`Video import failed: ${err.message}`, "error"); + } finally { + hideProgress(); + } +}; + +$("addImagesBtn").onclick = () => $("imageFiles").click(); +$("addVideoBtn").onclick = () => $("videoFile").click(); +$("imageFiles").onchange = (e) => { + uploadImages(e.target.files); + e.target.value = ""; +}; +$("videoFile").onchange = (e) => { + if (e.target.files[0]) askVideoStride(e.target.files[0]); + e.target.value = ""; +}; + +/* drag & drop */ + +let dragDepth = 0; +stage.addEventListener("dragenter", (e) => { + e.preventDefault(); + dragDepth += 1; + $("dropHint").classList.remove("hidden"); +}); +stage.addEventListener("dragleave", () => { + dragDepth = Math.max(0, dragDepth - 1); + if (dragDepth === 0) $("dropHint").classList.add("hidden"); +}); +stage.addEventListener("dragover", (e) => e.preventDefault()); +stage.addEventListener("drop", (e) => { + e.preventDefault(); + dragDepth = 0; + $("dropHint").classList.add("hidden"); + const files = Array.from(e.dataTransfer.files); + const videos = files.filter((f) => f.type.startsWith("video/")); + const images = files.filter((f) => f.type.startsWith("image/")); + if (images.length) uploadImages(images); + if (videos.length) askVideoStride(videos[0]); +}); + +/* ---------- canvas: view ---------- */ + +function resizeCanvas() { + const dpr = window.devicePixelRatio || 1; + const rect = stage.getBoundingClientRect(); + canvas.width = Math.round(rect.width * dpr); + canvas.height = Math.round(rect.height * dpr); + ctx.setTransform(dpr, 0, 0, dpr, 0, 0); + if (state.fitted) fitView(); + else render(); +} + +function fitView() { + const info = currentImage(); + if (!info) { + render(); + return; + } + const rect = stage.getBoundingClientRect(); + const pad = 24; + const scale = Math.min( + (rect.width - pad * 2) / info.width, + (rect.height - pad * 2) / info.height, + 1, + ); + state.view.scale = Math.max(scale, MIN_SCALE); + state.view.x = (rect.width - info.width * state.view.scale) / 2; + state.view.y = (rect.height - info.height * state.view.scale) / 2; + state.fitted = true; + updateZoomText(); + render(); +} + +function setZoom(newScale, cx, cy) { + const scale = clamp(newScale, MIN_SCALE, MAX_SCALE); + const before = screenToImage(cx, cy); + state.view.scale = scale; + state.view.x = cx - before.x * scale; + state.view.y = cy - before.y * scale; + state.fitted = false; + updateZoomText(); + render(); +} + +function updateZoomText() { + $("zoomText").textContent = `${Math.round(state.view.scale * 100)}%`; +} + +$("fitBtn").onclick = fitView; +$("zoomInBtn").onclick = () => { + const rect = stage.getBoundingClientRect(); + setZoom(state.view.scale * 1.25, rect.width / 2, rect.height / 2); +}; +$("zoomOutBtn").onclick = () => { + const rect = stage.getBoundingClientRect(); + setZoom(state.view.scale / 1.25, rect.width / 2, rect.height / 2); +}; + +canvas.addEventListener( + "wheel", + (e) => { + e.preventDefault(); + const factor = Math.exp(-e.deltaY * 0.0015); + setZoom(state.view.scale * factor, e.offsetX, e.offsetY); + }, + { passive: false }, +); + +/* ---------- canvas: hit testing ---------- */ + +function handlePositions(box) { + const s = state.view.scale; + const x = box.x * s + state.view.x; + const y = box.y * s + state.view.y; + const w = box.w * s; + const h = box.h * s; + return [ + { x, y, cursor: "nwse-resize", dx: -1, dy: -1 }, + { x: x + w / 2, y, cursor: "ns-resize", dx: 0, dy: -1 }, + { x: x + w, y, cursor: "nesw-resize", dx: 1, dy: -1 }, + { x: x + w, y: y + h / 2, cursor: "ew-resize", dx: 1, dy: 0 }, + { x: x + w, y: y + h, cursor: "nwse-resize", dx: 1, dy: 1 }, + { x: x + w / 2, y: y + h, cursor: "ns-resize", dx: 0, dy: 1 }, + { x, y: y + h, cursor: "nesw-resize", dx: -1, dy: 1 }, + { x, y: y + h / 2, cursor: "ew-resize", dx: -1, dy: 0 }, + ]; +} + +function hitTest(sx, sy) { + const boxes = currentBoxes(); + // Handles of the selected box take priority. + if (state.selection >= 0 && boxes[state.selection]) { + const handles = handlePositions(boxes[state.selection]); + for (let h = 0; h < handles.length; h++) { + if ( + Math.abs(sx - handles[h].x) <= HANDLE_HIT && + Math.abs(sy - handles[h].y) <= HANDLE_HIT + ) { + return { box: state.selection, handle: h }; + } + } + } + const pt = screenToImage(sx, sy); + let best = -1; + let bestArea = Infinity; + boxes.forEach((box, idx) => { + const inside = + pt.x >= box.x && + pt.x <= box.x + box.w && + pt.y >= box.y && + pt.y <= box.y + box.h; + const area = box.w * box.h; + if (inside && area < bestArea) { + best = idx; + bestArea = area; + } + }); + return { box: best, handle: -1 }; +} + +/* ---------- canvas: interactions ---------- */ + +canvas.addEventListener("pointerdown", (e) => { + if (!currentImage() && e.button === 0) return; + canvas.setPointerCapture(e.pointerId); + const info = currentImage(); + + if (e.button === 1 || (e.button === 0 && (spaceDown || !info))) { + drag = { + mode: "pan", + startX: e.offsetX, + startY: e.offsetY, + viewX: state.view.x, + viewY: state.view.y, + }; + return; + } + if (e.button !== 0) return; + + const hit = hitTest(e.offsetX, e.offsetY); + const boxes = currentBoxes(); + + if (hit.handle >= 0) { + const box = boxes[hit.box]; + drag = { + mode: "resize", + index: hit.box, + handle: hit.handle, + orig: { ...box }, + moved: false, + }; + return; + } + if (hit.box >= 0) { + state.selection = hit.box; + const pt = screenToImage(e.offsetX, e.offsetY); + const box = boxes[hit.box]; + drag = { + mode: "move", + index: hit.box, + offsetX: pt.x - box.x, + offsetY: pt.y - box.y, + moved: false, + }; + render(); + return; + } + + state.selection = -1; + if (state.activeLabel) { + const pt = screenToImage(e.offsetX, e.offsetY); + const x = clamp(pt.x, 0, info.width); + const y = clamp(pt.y, 0, info.height); + drag = { mode: "draw", startX: x, startY: y, rect: null }; + } else { + drag = { + mode: "pan", + startX: e.offsetX, + startY: e.offsetY, + viewX: state.view.x, + viewY: state.view.y, + }; + } + render(); +}); + +canvas.addEventListener("pointermove", (e) => { + pointer = { x: e.offsetX, y: e.offsetY, inside: true }; + const info = currentImage(); + + if (!drag) { + hover = info ? hitTest(e.offsetX, e.offsetY) : { box: -1, handle: -1 }; + updateCursor(); + render(); + return; + } + + if (drag.mode === "pan") { + state.view.x = drag.viewX + (e.offsetX - drag.startX); + state.view.y = drag.viewY + (e.offsetY - drag.startY); + state.fitted = false; + render(); + return; + } + + const pt = screenToImage(e.offsetX, e.offsetY); + const boxes = currentBoxes(); + + if (drag.mode === "draw") { + const x = clamp(pt.x, 0, info.width); + const y = clamp(pt.y, 0, info.height); + drag.rect = { + x: Math.min(drag.startX, x), + y: Math.min(drag.startY, y), + w: Math.abs(x - drag.startX), + h: Math.abs(y - drag.startY), + }; + } else if (drag.mode === "move") { + const box = boxes[drag.index]; + box.x = clamp(pt.x - drag.offsetX, 0, info.width - box.w); + box.y = clamp(pt.y - drag.offsetY, 0, info.height - box.h); + drag.moved = true; + } else if (drag.mode === "resize") { + resizeBox(boxes[drag.index], drag, pt, info); + drag.moved = true; + } + render(); +}); + +function resizeBox(box, dragState, pt, info) { + const orig = dragState.orig; + const handle = handlePositions(orig)[dragState.handle]; + let x1 = orig.x; + let y1 = orig.y; + let x2 = orig.x + orig.w; + let y2 = orig.y + orig.h; + const px = clamp(pt.x, 0, info.width); + const py = clamp(pt.y, 0, info.height); + if (handle.dx < 0) x1 = px; + if (handle.dx > 0) x2 = px; + if (handle.dy < 0) y1 = py; + if (handle.dy > 0) y2 = py; + box.x = Math.min(x1, x2); + box.y = Math.min(y1, y2); + box.w = Math.abs(x2 - x1); + box.h = Math.abs(y2 - y1); +} + +canvas.addEventListener("pointerup", (e) => { + if (!drag) return; + const info = currentImage(); + + if (drag.mode === "draw" && drag.rect && info) { + if (drag.rect.w >= MIN_BOX_SIZE && drag.rect.h >= MIN_BOX_SIZE) { + const boxes = currentBoxes(); + boxes.push({ label: state.activeLabel, ...drag.rect }); + state.selection = boxes.length - 1; + markChanged(); + renderLabels(); + renderImages(); + } + } else if ( + (drag.mode === "move" || drag.mode === "resize") && + drag.moved + ) { + markChanged(); + } + drag = null; + hover = info ? hitTest(e.offsetX, e.offsetY) : { box: -1, handle: -1 }; + updateCursor(); + render(); +}); + +canvas.addEventListener("pointerleave", () => { + pointer.inside = false; + hover = { box: -1, handle: -1 }; + render(); +}); + +canvas.addEventListener("contextmenu", (e) => { + e.preventDefault(); + if (!currentImage()) return; + const hit = hitTest(e.offsetX, e.offsetY); + if (hit.box >= 0) deleteBox(hit.box); +}); + +function deleteBox(index) { + const boxes = currentBoxes(); + boxes.splice(index, 1); + if (state.selection === index) state.selection = -1; + else if (state.selection > index) state.selection -= 1; + markChanged(); + renderLabels(); + renderImages(); + render(); +} + +function updateCursor() { + if (spaceDown || (drag && drag.mode === "pan")) { + canvas.style.cursor = "grab"; + } else if (hover.handle >= 0) { + const boxes = currentBoxes(); + canvas.style.cursor = handlePositions(boxes[hover.box])[ + hover.handle + ].cursor; + } else if (hover.box >= 0) { + canvas.style.cursor = "move"; + } else if (state.activeLabel && currentImage()) { + canvas.style.cursor = "crosshair"; + } else { + canvas.style.cursor = "default"; + } +} + +/* ---------- keyboard ---------- */ + +document.addEventListener("keydown", (e) => { + const tag = document.activeElement && document.activeElement.tagName; + if (tag === "INPUT" || tag === "TEXTAREA") return; + if (document.querySelector("dialog[open]")) return; + + if (e.code === "Space") { + spaceDown = true; + updateCursor(); + e.preventDefault(); + } else if (e.key === "Escape") { + if (drag && drag.mode === "draw") drag = null; + state.selection = -1; + render(); + } else if (e.key === "Delete" || e.key === "Backspace") { + if (state.selection >= 0) { + deleteBox(state.selection); + e.preventDefault(); + } + } else if (e.key === "ArrowLeft") { + $("prevBtn").click(); + } else if (e.key === "ArrowRight") { + $("nextBtn").click(); + } +}); + +document.addEventListener("keyup", (e) => { + if (e.code === "Space") { + spaceDown = false; + updateCursor(); + } +}); + +/* ---------- rendering ---------- */ + +function render() { + const rect = stage.getBoundingClientRect(); + ctx.clearRect(0, 0, rect.width, rect.height); + const info = currentImage(); + if (!info) return; + + const img = imageCache.get(info.name); + const { scale, x: ox, y: oy } = state.view; + + if (img && img.complete && img.naturalWidth) { + ctx.imageSmoothingEnabled = scale < 4; + ctx.drawImage(img, ox, oy, info.width * scale, info.height * scale); + } + ctx.strokeStyle = "rgba(255,255,255,0.25)"; + ctx.lineWidth = 1; + ctx.strokeRect(ox, oy, info.width * scale, info.height * scale); + + const boxes = currentBoxes(); + boxes.forEach((box, idx) => { + drawBox(box, idx === state.selection, idx === hover.box); + }); + + if (drag && drag.mode === "draw" && drag.rect) { + drawBox({ label: state.activeLabel, ...drag.rect }, false, false); + } + + // Crosshair guides while drawing is possible. + const drawing = drag && drag.mode === "draw"; + const idle = !drag && hover.box < 0 && state.activeLabel && !spaceDown; + if (pointer.inside && (drawing || idle)) { + ctx.save(); + ctx.strokeStyle = "rgba(230,233,239,0.35)"; + ctx.lineWidth = 1; + ctx.setLineDash([5, 5]); + ctx.beginPath(); + ctx.moveTo(pointer.x, 0); + ctx.lineTo(pointer.x, rect.height); + ctx.moveTo(0, pointer.y); + ctx.lineTo(rect.width, pointer.y); + ctx.stroke(); + ctx.restore(); + } +} + +function drawBox(box, selected, hovered) { + const { scale, x: ox, y: oy } = state.view; + const x = box.x * scale + ox; + const y = box.y * scale + oy; + const w = box.w * scale; + const h = box.h * scale; + const color = labelColor(box.label); + + ctx.save(); + ctx.fillStyle = hexToRgba(color, selected || hovered ? 0.22 : 0.12); + ctx.fillRect(x, y, w, h); + ctx.strokeStyle = color; + ctx.lineWidth = selected ? 2.5 : 2; + ctx.setLineDash(selected ? [] : [6, 4]); + ctx.strokeRect(x, y, w, h); + ctx.setLineDash([]); + + // Label tag. + const text = box.label || ""; + ctx.font = "11px " + getComputedStyle(document.body).fontFamily; + const tw = ctx.measureText(text).width + 10; + const ty = y - 17 < 0 ? y : y - 17; + ctx.fillStyle = color; + ctx.fillRect(x, ty, tw, 17); + ctx.fillStyle = "#fff"; + ctx.fillText(text, x + 5, ty + 12); + + if (selected) { + for (const handle of handlePositions(box)) { + ctx.fillStyle = "#fff"; + ctx.strokeStyle = color; + ctx.lineWidth = 1.5; + ctx.fillRect( + handle.x - HANDLE_SIZE / 2, + handle.y - HANDLE_SIZE / 2, + HANDLE_SIZE, + HANDLE_SIZE, + ); + ctx.strokeRect( + handle.x - HANDLE_SIZE / 2, + handle.y - HANDLE_SIZE / 2, + HANDLE_SIZE, + HANDLE_SIZE, + ); + } + } + ctx.restore(); +} + +function hexToRgba(hex, alpha) { + const value = parseInt(hex.slice(1), 16); + const r = (value >> 16) & 255; + const g = (value >> 8) & 255; + const b = value & 255; + return `rgba(${r},${g},${b},${alpha})`; +} + +/* ---------- dialogs ---------- */ + +for (const dialog of document.querySelectorAll("dialog")) { + const closeBtn = dialog.querySelector("[data-close]"); + if (closeBtn) closeBtn.onclick = () => dialog.close(); +} + +$("exportBtn").onclick = () => $("exportDialog").showModal(); +$("synthBtn").onclick = () => $("synthDialog").showModal(); + +$("exportForm").onsubmit = async (e) => { + e.preventDefault(); + const form = e.target.elements; + const submitBtn = e.target.querySelector("button[type=submit]"); + submitBtn.disabled = true; + await syncSession(); + showProgress("Exporting dataset…"); + startProgressPoll("Exporting dataset"); + try { + const result = await api("/api/export", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + name: form.name.value.trim() || "yolo_dataset", + train_split: Number(form.train.value) / 100, + shuffle: form.shuffle.checked, + seed: form.seed.value === "" ? null : Number(form.seed.value), + output_dir: form.output.value.trim() || null, + }), + }); + $("exportDialog").close(); + toast( + `Dataset exported (${result.num_train} train / ` + + `${result.num_val} val).`, + "ok", + result.yaml, + 12000, + ); + } catch (err) { + toast(`Export failed: ${err.message}`, "error"); + } finally { + hideProgress(); + submitBtn.disabled = false; + } +}; + +$("synthForm").onsubmit = async (e) => { + e.preventDefault(); + const form = e.target.elements; + const submitBtn = e.target.querySelector("button[type=submit]"); + submitBtn.disabled = true; + await syncSession(); + showProgress("Synthesizing dataset…"); + startProgressPoll("Synthesizing dataset"); + try { + const result = await api("/api/synthesize", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + name: form.name.value.trim() || "synt_dataset", + num_images: Number(form.num.value), + width: Number(form.width.value), + height: Number(form.height.value), + per_image: Number(form.per.value), + train_split: Number(form.train.value) / 100, + scale_min: Number(form.smin.value), + scale_max: Number(form.smax.value), + background: form.background.value, + output_dir: form.output.value.trim() || null, + }), + }); + $("synthDialog").close(); + toast( + `Synthetic dataset created (${result.num_train} train / ` + + `${result.num_val} val).`, + "ok", + result.yaml, + 12000, + ); + } catch (err) { + toast(`Synthesis failed: ${err.message}`, "error"); + } finally { + hideProgress(); + submitBtn.disabled = false; + } +}; + +/* ---------- session save / load ---------- */ + +$("saveSessionBtn").onclick = () => { + quitAfterSave = false; + openSaveDialog(); +}; + +function openSaveDialog() { + const input = $("saveForm").elements.path; + if (sessionPath && !input.value) input.value = sessionPath; + $("saveDialog").showModal(); +} + +$("saveForm").onsubmit = async (e) => { + e.preventDefault(); + const path = e.target.elements.path.value.trim(); + if (!path) return; + try { + const result = await api("/api/session", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + path, + labels: state.labels, + annotations: state.annotations, + }), + }); + sessionPath = result.path; + dirty = false; + setSaveStatus(); + $("saveDialog").close(); + toast("Session saved.", "ok", result.path); + if (quitAfterSave) { + quitAfterSave = false; + await shutdownServer(); + } + } catch (err) { + toast(`Save failed: ${err.message}`, "error"); + } +}; + +$("loadSessionBtn").onclick = () => { + const input = $("loadForm").elements.path; + if (sessionPath && !input.value) input.value = sessionPath; + $("loadDialog").showModal(); +}; + +$("loadForm").onsubmit = async (e) => { + e.preventDefault(); + const path = e.target.elements.path.value.trim(); + if (!path) return; + if ( + dirty && + !confirm("Loading a session discards unsaved changes. Continue?") + ) { + return; + } + try { + const session = await api("/api/session/load", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ path }), + }); + state.labels = session.labels || []; + state.annotations = session.annotations || {}; + state.activeLabel = null; + state.selection = -1; + sessionPath = path; + dirty = false; + setSaveStatus(); + $("loadDialog").close(); + renderLabels(); + renderImages(); + render(); + toast("Session loaded.", "ok", path); + } catch (err) { + toast(`Load failed: ${err.message}`, "error"); + } +}; + +/* ---------- quit ---------- */ + +async function shutdownServer() { + dirty = false; // suppress the beforeunload warning + try { + await api("/api/shutdown", { method: "POST" }); + } catch { + /* server is going down */ + } + document.body.innerHTML = + '

Server stopped.

' + + "

You can close this tab.

"; +} + +$("quitBtn").onclick = () => { + if (dirty) $("quitDialog").showModal(); + else shutdownServer(); +}; + +$("quitDiscardBtn").onclick = () => { + $("quitDialog").close(); + shutdownServer(); +}; + +$("quitSaveBtn").onclick = () => { + $("quitDialog").close(); + quitAfterSave = true; + openSaveDialog(); +}; + +/* ---------- init ---------- */ + +async function init() { + resizeCanvas(); + try { + const data = await api("/api/state"); + state.workspace = data.workspace; + state.images = data.images; + state.labels = data.labels; + state.annotations = data.annotations; + sessionPath = data.session_path; + dirty = Boolean(data.dirty); + $("workspacePath").textContent = data.workspace; + setSaveStatus(); + renderLabels(); + renderImages(); + if (state.images.length) selectImage(0, true); + } catch (err) { + toast(`Could not load session: ${err.message}`, "error"); + } +} + +window.addEventListener("resize", resizeCanvas); +new ResizeObserver(resizeCanvas).observe(stage); +init(); diff --git a/src/dynsight/_internal/vision/label_tool/index.html b/src/dynsight/_internal/vision/label_tool/index.html index d7e05bfd..e5e0c36d 100644 --- a/src/dynsight/_internal/vision/label_tool/index.html +++ b/src/dynsight/_internal/vision/label_tool/index.html @@ -3,55 +3,281 @@ - dynsight labeling tool + dynsight · label tool + -
-

Labels

- - -
-
+
+ -
-
- - - - - - - - - - -
-
-
- -
-
-
-
+
+ + 0 / 0 +
+
+ + 100% + +
+ +
+ + + + + + +
+ +
+ + +
+ +
+

No images yet

+

+ Use + Images / + Video or drop files + anywhere on this area. +

+
+ + +
- - +
+ + + + + draw: drag · select: click · move/resize: drag + box · delete: ⌫ or right-click · pan: + space/middle drag · zoom: wheel + +
+ + +
+

Export YOLO dataset

+ +
+ + +
+ + +
+ + +
+
+ + + +
+

Synthesize dataset

+

+ Annotated crops are pasted at random positions onto uniform + backgrounds. +

+ +
+ + +
+
+ + +
+
+ + +
+
+ + +
+ +
+ + +
+
+ + + +
+

Save session

+

+ Labels and boxes are written to a JSON file that can be + loaded back later. +

+ +
+ + +
+
+ + + +
+

Load session

+

+ Loading a session replaces the current labels and boxes. +

+ +
+ + +
+
+ + + +
+

Unsaved session

+

+ The current session has not been saved to a file. Do you + want to save it before quitting? +

+
+ + + +
+
+ + + +
+

Import video frames

+

+ +
+ + +
+
+ + +
+ + + + + diff --git a/src/dynsight/_internal/vision/label_tool/logo.png b/src/dynsight/_internal/vision/label_tool/logo.png new file mode 100644 index 00000000..815d096d Binary files /dev/null and b/src/dynsight/_internal/vision/label_tool/logo.png differ diff --git a/src/dynsight/_internal/vision/label_tool/script.js b/src/dynsight/_internal/vision/label_tool/script.js deleted file mode 100644 index 8174e41d..00000000 --- a/src/dynsight/_internal/vision/label_tool/script.js +++ /dev/null @@ -1,577 +0,0 @@ -// script.js - -const imageInput = document.getElementById("imageInput"); -const imageDisplay = document.getElementById("imageDisplay"); -const imageContainer = document.getElementById("imageContainer"); -const imageWrapper = document.getElementById("imageWrapper"); -const labelList = document.getElementById("labelList"); -const addLabelBtn = document.getElementById("addLabelBtn"); -const newLabelInput = document.getElementById("newLabelInput"); -const clearLastBtn = document.getElementById("clearLast"); -const clearAllBtn = document.getElementById("clearAll"); -const exportBtn = document.getElementById("exportYolo"); -const exportAllBtn = document.getElementById("exportAll"); -const synthBtn = document.getElementById("synthesize"); -const nextImageBtn = document.getElementById("nextImage"); -const prevImageBtn = document.getElementById("prevImage"); -const verticalLine = document.getElementById("verticalLine"); -const horizontalLine = document.getElementById("horizontalLine"); -const zoomSlider = document.getElementById("zoomSlider"); - -let zoomLevel = 1; -let baseZoom = 1; -let naturalWidth = 0; -let naturalHeight = 0; - -const overlay = document.getElementById("overlay"); - -verticalLine.style.display = "none"; -horizontalLine.style.display = "none"; - -imageContainer.onmouseenter = () => { - verticalLine.style.display = "block"; - horizontalLine.style.display = "block"; -}; - -imageContainer.onmouseleave = () => { - verticalLine.style.display = "none"; - horizontalLine.style.display = "none"; -}; - -let currentLabel = null; -const labelColors = {}; -let isDrawing = false; -let startX, - startY, - box = null; - -let images = []; -let currentIndex = 0; -const annotations = {}; // imageName -> [boxData] - -function getRandomColor() { - const hue = Math.floor(Math.random() * 360); - return `hsl(${hue}, 90%, 50%)`; -} - -function setActiveLabel(item) { - document - .querySelectorAll(".label-item") - .forEach((i) => i.classList.remove("active")); - item.classList.add("active"); - currentLabel = item.textContent; -} - -function createLabelItem(text) { - const item = document.createElement("div"); - item.className = "label-item"; - item.textContent = text; - labelColors[text] = labelColors[text] || getRandomColor(); - item.style.backgroundColor = labelColors[text]; - item.style.color = "#fff"; - item.addEventListener("click", () => setActiveLabel(item)); - labelList.appendChild(item); -} - -addLabelBtn.onclick = () => { - const label = newLabelInput.value.trim(); - if (label && !labelColors[label]) { - createLabelItem(label); - newLabelInput.value = ""; - } -}; - -imageInput.onchange = (e) => { - images = Array.from(e.target.files); - currentIndex = 0; - loadImage(currentIndex); -}; - -function loadImage(index) { - if (!images[index]) return; - const url = URL.createObjectURL(images[index]); - imageDisplay.onload = () => { - const iw = imageDisplay.naturalWidth; - const ih = imageDisplay.naturalHeight; - imageDisplay.style.width = `${iw}px`; - imageDisplay.style.height = `${ih}px`; - naturalWidth = iw; - naturalHeight = ih; - baseZoom = Math.min( - imageContainer.clientWidth / iw, - imageContainer.clientHeight / ih, - 1, - ); - zoomLevel = baseZoom; - zoomSlider.value = 100; - updateTransform(); - const name = images[index].name; - if (!annotations[name]) annotations[name] = []; - clearBoxes(); - annotations[name].forEach(addBoxFromData); - }; - imageDisplay.src = url; -} - -zoomSlider.oninput = (e) => { - zoomLevel = (e.target.value / 100) * baseZoom; - updateTransform(); - clearBoxes(); - annotations[images[currentIndex].name].forEach(addBoxFromData); -}; - -function clearBoxes() { - overlay.innerHTML = ""; -} - -function updateTransform() { - const w = naturalWidth * zoomLevel; - const h = naturalHeight * zoomLevel; - imageDisplay.style.width = `${w}px`; - imageDisplay.style.height = `${h}px`; - imageWrapper.style.width = `${w}px`; - imageWrapper.style.height = `${h}px`; - overlay.style.width = `${w}px`; - overlay.style.height = `${h}px`; -} - -function addBoxFromData(data) { - const box = document.createElement("div"); - box.className = "bounding-box"; - box.style.left = `${data.left * zoomLevel}px`; - box.style.top = `${data.top * zoomLevel}px`; - box.style.width = `${data.width * zoomLevel}px`; - box.style.height = `${data.height * zoomLevel}px`; - box.style.border = `2px dashed ${labelColors[data.label]}`; - box.style.backgroundColor = labelColors[data.label] - .replace("hsl", "hsla") - .replace(")", ", 0.1)"); - - const tag = document.createElement("div"); - tag.className = "label-tag"; - tag.textContent = data.label; - tag.style.backgroundColor = labelColors[data.label]; - box.appendChild(tag); - - overlay.appendChild(box); -} - -imageContainer.onmousedown = (e) => { - if (e.button !== 0) { - return; - } - - if (!currentLabel || !images[currentIndex]) return; - - const rect = imageDisplay.getBoundingClientRect(); - startX = (e.clientX - rect.left) / zoomLevel; - startY = (e.clientY - rect.top) / zoomLevel; - - // Ignore clicks started outside the image boundaries - if ( - startX < 0 || - startY < 0 || - startX > naturalWidth || - startY > naturalHeight - ) { - isDrawing = false; - return; - } - - box = document.createElement("div"); - box.className = "bounding-box"; - box.style.left = `${startX * zoomLevel}px`; - box.style.top = `${startY * zoomLevel}px`; - box.style.border = `2px dashed ${labelColors[currentLabel]}`; - box.style.backgroundColor = labelColors[currentLabel] - .replace("hsl", "hsla") - .replace(")", ", 0.1)"); - - const tag = document.createElement("div"); - tag.className = "label-tag"; - tag.textContent = currentLabel; - tag.style.backgroundColor = labelColors[currentLabel]; - box.appendChild(tag); - - overlay.appendChild(box); - isDrawing = true; -}; - -imageContainer.onmousemove = (e) => { - const imgRect = imageDisplay.getBoundingClientRect(); - const containerRect = imageContainer.getBoundingClientRect(); - - const currX = (e.clientX - imgRect.left) / zoomLevel; - const currY = (e.clientY - imgRect.top) / zoomLevel; - const clampedX = Math.max(0, Math.min(naturalWidth, currX)); - const clampedY = Math.max(0, Math.min(naturalHeight, currY)); - - verticalLine.style.left = `${ - e.clientX - containerRect.left + imageContainer.scrollLeft - }px`; - horizontalLine.style.top = `${ - e.clientY - containerRect.top + imageContainer.scrollTop - }px`; - - if (!isDrawing || !box) return; - - box.style.left = `${Math.min(clampedX, startX) * zoomLevel}px`; - box.style.top = `${Math.min(clampedY, startY) * zoomLevel}px`; - box.style.width = `${Math.abs(clampedX - startX) * zoomLevel}px`; - box.style.height = `${Math.abs(clampedY - startY) * zoomLevel}px`; -}; - -imageContainer.onmouseup = (e) => { - if (!isDrawing || !box) return; - - const imgRect = imageDisplay.getBoundingClientRect(); - const endX = (e.clientX - imgRect.left) / zoomLevel; - const endY = (e.clientY - imgRect.top) / zoomLevel; - const clampedX = Math.max(0, Math.min(naturalWidth, endX)); - const clampedY = Math.max(0, Math.min(naturalHeight, endY)); - - const left = Math.min(startX, clampedX); - const top = Math.min(startY, clampedY); - const width = Math.abs(clampedX - startX); - const height = Math.abs(clampedY - startY); - - annotations[images[currentIndex].name].push({ - label: currentLabel, - left, - top, - width, - height, - }); - - box = null; - isDrawing = false; -}; - -clearLastBtn.onclick = () => { - const ann = annotations[images[currentIndex].name]; - if (ann.length > 0) { - ann.pop(); - clearBoxes(); - ann.forEach(addBoxFromData); - } -}; - -clearAllBtn.onclick = () => { - annotations[images[currentIndex].name] = []; - clearBoxes(); -}; - -prevImageBtn.onclick = () => { - if (currentIndex > 0) { - currentIndex--; - loadImage(currentIndex); - } -}; - -nextImageBtn.onclick = () => { - if (currentIndex < images.length - 1) { - currentIndex++; - loadImage(currentIndex); - } -}; - -exportBtn.onclick = () => { - const img = images[currentIndex]; - if (!img) return; - const iw = imageDisplay.naturalWidth; - const ih = imageDisplay.naturalHeight; - const annots = annotations[img.name] || []; - const labelMap = {}; - let nextId = 0; - let txt = ""; - annots.forEach((obj) => { - if (!(obj.label in labelMap)) labelMap[obj.label] = nextId++; - const cx = (obj.left + obj.width / 2) / iw; - const cy = (obj.top + obj.height / 2) / ih; - const w = obj.width / iw; - const h = obj.height / ih; - txt += - labelMap[obj.label] + - " " + - cx.toFixed(6) + - " " + - cy.toFixed(6) + - " " + - w.toFixed(6) + - " " + - h.toFixed(6) + - "\n"; - }); - const blob = new Blob([txt], { type: "text/plain" }); - const a = document.createElement("a"); - a.href = URL.createObjectURL(blob); - a.download = img.name.replace(/\.[^/.]+$/, "") + ".txt"; - a.click(); - URL.revokeObjectURL(a.href); -}; - -exportAllBtn.onclick = async () => { - if (images.length === 0) { - alert("No images uploaded."); - return; - } - let trainPercent = parseFloat( - prompt("Percentage of images for training?", "80"), - ); - if ( - Number.isNaN(trainPercent) || - trainPercent <= 0 || - trainPercent >= 100 - ) { - trainPercent = 80; - } - const numTrain = Math.floor(images.length * (trainPercent / 100)); - const zip = new JSZip(); - const datasetName = "yolo_dataset"; - const imgTrain = zip.folder("images/train"); - const imgVal = zip.folder("images/val"); - const lblTrain = zip.folder("labels/train"); - const lblVal = zip.folder("labels/val"); - const labelMap = {}; - let nextId = 0; - for (let i = 0; i < images.length; i++) { - const image = images[i]; - const name = image.name; - const imgData = await image.arrayBuffer(); - const imgFolder = i < numTrain ? imgTrain : imgVal; - const lblFolder = i < numTrain ? lblTrain : lblVal; - imgFolder.file(name, imgData); - - const img = new Image(); - const url = URL.createObjectURL(image); - img.src = url; - await new Promise((resolve) => (img.onload = resolve)); - const iw = img.naturalWidth; - const ih = img.naturalHeight; - const annots = annotations[name] || []; - let txt = ""; - annots.forEach((obj) => { - if (!(obj.label in labelMap)) labelMap[obj.label] = nextId++; - const cx = (obj.left + obj.width / 2) / iw; - const cy = (obj.top + obj.height / 2) / ih; - const w = obj.width / iw; - const h = obj.height / ih; - txt += - labelMap[obj.label] + - " " + - cx.toFixed(6) + - " " + - cy.toFixed(6) + - " " + - w.toFixed(6) + - " " + - h.toFixed(6) + - "\n"; - }); - const labelFileName = name.replace(/\.[^/.]+$/, "") + ".txt"; - lblFolder.file(labelFileName, txt); - } - - const names = Object.keys(labelMap); - const yaml = [ - `path: ${datasetName}`, - "train: images/train", - "val: images/val", - `nc: ${names.length}`, - `names: [${names.map((n) => `'${n}'`).join(", ")}]`, - ].join("\n"); - zip.file("dataset.yaml", yaml); - - const content = await zip.generateAsync({ type: "blob" }); - const a = document.createElement("a"); - a.href = URL.createObjectURL(content); - a.download = `${datasetName}.zip`; - a.click(); - URL.revokeObjectURL(a.href); -}; - -synthBtn.onclick = async () => { - if (images.length === 0) { - alert("No images uploaded."); - return; - } - - const numImages = parseInt( - prompt("Number of synthetic images to generate?", "10"), - 10, - ); - const width = parseInt(prompt("Image width?", "640"), 10); - const height = parseInt(prompt("Image height?", "640"), 10); - const requestedPerImage = parseInt( - prompt("Number of objects per image?", "10"), - 10, - ); - if ( - !numImages || - Number.isNaN(numImages) || - !width || - Number.isNaN(width) || - !height || - Number.isNaN(height) || - !requestedPerImage || - Number.isNaN(requestedPerImage) || - requestedPerImage < 1 - ) { - alert("Invalid parameters."); - return; - } - - const crops = []; - const datasetName = "synt_dataset"; - const labelMap = {}; - let nextId = 0; - for (const file of images) { - const ann = annotations[file.name] || []; - for (const c of ann) { - crops.push({ file, ...c }); - if (!(c.label in labelMap)) labelMap[c.label] = nextId++; - } - } - - if (crops.length === 0) { - alert("No label found."); - return; - } - - async function loadImage(file) { - return await new Promise((resolve) => { - const img = new Image(); - img.src = URL.createObjectURL(file); - img.onload = () => { - URL.revokeObjectURL(img.src); - resolve(img); - }; - }); - } - - function overlaps(x, y, w, h, boxes) { - return boxes.some((b) => { - return !( - x + w <= b.x || - x >= b.x + b.w || - y + h <= b.y || - y >= b.y + b.h - ); - }); - } - - async function createCollage() { - const canvas = document.createElement("canvas"); - canvas.width = width; - canvas.height = height; - const ctx = canvas.getContext("2d"); - ctx.fillStyle = "white"; - ctx.fillRect(0, 0, width, height); - - const placed = []; - for (let i = 0; i < requestedPerImage; i++) { - const crop = crops[Math.floor(Math.random() * crops.length)]; - const img = await loadImage(crop.file); - const c = document.createElement("canvas"); - c.width = crop.width; - c.height = crop.height; - c.getContext("2d").drawImage( - img, - crop.left, - crop.top, - crop.width, - crop.height, - 0, - 0, - crop.width, - crop.height, - ); - - const scale = 1 - const w = crop.width * scale; - const h = crop.height * scale; - - let x, y; - let tries = 0; - do { - x = Math.random() * (width - w); - y = Math.random() * (height - h); - tries += 1; - } while (tries < 50 && overlaps(x, y, w, h, placed)); - - if (tries === 50) continue; - - ctx.drawImage(c, 0, 0, crop.width, crop.height, x, y, w, h); - placed.push({ label: crop.label, x, y, w, h }); - } - - let txt = ""; - placed.forEach((p) => { - const cls = labelMap[p.label]; - const cx = (p.x + p.w / 2) / width; - const cy = (p.y + p.h / 2) / height; - const ww = p.w / width; - const hh = p.h / height; - txt += `${cls} ${cx.toFixed(6)} ${cy.toFixed(6)} ${ww.toFixed( - 6, - )} ${hh.toFixed(6)}\n`; - }); - - const blob = await new Promise((resolve) => - canvas.toBlob((b) => resolve(b), "image/jpeg"), - ); - return { blob, txt }; - } - - const zip = new JSZip(); - const imgTrain = zip.folder("images/train"); - const imgVal = zip.folder("images/val"); - const lblTrain = zip.folder("labels/train"); - const lblVal = zip.folder("labels/val"); - - const numTrain = Math.floor(numImages * 0.8); - - for (let i = 0; i < numImages; i++) { - const { blob, txt } = await createCollage(); - const imgName = `synt_${i}.jpg`; - const txtName = `synt_${i}.txt`; - if (i < numTrain) { - imgTrain.file(imgName, blob); - lblTrain.file(txtName, txt); - } else { - imgVal.file(imgName, blob); - lblVal.file(txtName, txt); - } - } - - const names = Object.keys(labelMap); - const yaml = [ - `path: ${datasetName}`, - "train: images/train", - "val: images/val", - `nc: ${names.length}`, - `names: [${names.map((n) => `'${n}'`).join(", ")}]`, - ].join("\n"); - zip.file("dataset.yaml", yaml); - - const content = await zip.generateAsync({ type: "blob" }); - const a = document.createElement("a"); - a.href = URL.createObjectURL(content); - a.download = `${datasetName}.zip`; - a.click(); - URL.revokeObjectURL(a.href); -}; - -let navigatingAway = false; -document.addEventListener("click", (e) => { - const link = e.target.closest("a"); - if (link && link.href) { - navigatingAway = true; - } -}); - -window.addEventListener("pagehide", () => { - if (!navigatingAway) { - navigator.sendBeacon("/shutdown"); - } -}); diff --git a/src/dynsight/_internal/vision/label_tool/styles.css b/src/dynsight/_internal/vision/label_tool/styles.css index fffb761f..4ec6bcd1 100644 --- a/src/dynsight/_internal/vision/label_tool/styles.css +++ b/src/dynsight/_internal/vision/label_tool/styles.css @@ -1,145 +1,568 @@ +:root { + --bg: #0f1115; + --bg-panel: #161a21; + --bg-raised: #1d232d; + --border: #2a313d; + --text: #e6e9ef; + --text-dim: #8b94a3; + --accent: #4f8cff; + --accent-hover: #669aff; + --danger: #f0506e; + --ok: #34c98e; + --radius: 8px; + --font: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, + "Helvetica Neue", Arial, sans-serif; + --mono: "SF Mono", SFMono-Regular, Menlo, Consolas, monospace; +} + * { - user-select: none; + box-sizing: border-box; } + +html, body { + height: 100%; margin: 0; - font-family: Arial, sans-serif; +} + +body { + display: flex; + flex-direction: column; + background: var(--bg); + color: var(--text); + font-family: var(--font); + font-size: 13px; + overflow: hidden; +} + +/* ---------- top bar ---------- */ + +#topbar { + display: flex; + align-items: center; + flex-wrap: wrap; + gap: 8px 12px; + padding: 8px 14px; + background: var(--bg-panel); + border-bottom: 1px solid var(--border); + flex: 0 0 auto; +} + +#logo { + height: 28px; + width: auto; + margin-right: 8px; + /* The logo has dark strokes: lift it slightly on the dark theme. */ + filter: drop-shadow(0 0 1px rgba(255, 255, 255, 0.25)); + user-select: none; + -webkit-user-drag: none; +} + +.group { + display: flex; + align-items: center; + gap: 4px; + background: var(--bg-raised); + border: 1px solid var(--border); + border-radius: var(--radius); + padding: 2px; +} + +.counter { + min-width: 52px; + text-align: center; + color: var(--text-dim); + font-variant-numeric: tabular-nums; +} + +.spacer { + flex: 1; +} + +/* ---------- buttons ---------- */ + +button { + font-family: inherit; + font-size: 13px; + color: var(--text); + background: transparent; + border: none; + border-radius: 6px; + padding: 6px 10px; + cursor: pointer; + white-space: nowrap; + transition: background 0.12s, border-color 0.12s, color 0.12s; +} + +button:disabled { + opacity: 0.45; + cursor: not-allowed; +} + +.icon-btn { + width: 28px; + height: 26px; + padding: 0; + line-height: 1; +} + +.icon-btn:hover:not(:disabled), +.ghost-btn:hover:not(:disabled) { + background: rgba(255, 255, 255, 0.07); +} + +.ghost-btn { + border: 1px solid var(--border); + background: var(--bg-raised); +} + +.primary-btn { + background: var(--accent); + color: #fff; + font-weight: 600; +} + +.primary-btn:hover:not(:disabled) { + background: var(--accent-hover); +} + +.danger-btn { + border: 1px solid transparent; + color: var(--danger); +} + +.danger-btn:hover:not(:disabled) { + border-color: var(--danger); + background: rgba(240, 80, 110, 0.1); +} + +/* ---------- layout ---------- */ + +#layout { display: flex; - height: 100vh; + flex: 1 1 auto; + min-height: 0; } -.sidebar { + +#sidebar { width: 250px; - background: #f4f4f4; - border-right: 1px solid #ccc; - padding: 20px; - box-sizing: border-box; - flex-shrink: 0; + flex: 0 0 auto; + display: flex; + flex-direction: column; + background: var(--bg-panel); + border-right: 1px solid var(--border); + overflow: hidden; +} + +.panel { + display: flex; + flex-direction: column; + padding: 12px; + border-bottom: 1px solid var(--border); + min-height: 0; + flex: 0 0 auto; +} + +.panel.grow { + flex: 1 1 auto; + border-bottom: none; } -.main { + +.panel h2 { + margin: 0 0 8px; + font-size: 11px; + font-weight: 600; + text-transform: uppercase; + letter-spacing: 0.08em; + color: var(--text-dim); +} + +#labelForm { + display: flex; + gap: 6px; + margin-bottom: 8px; +} + +#labelForm input { flex: 1; min-width: 0; +} + +input[type="text"], +input[type="number"] { + background: var(--bg); + border: 1px solid var(--border); + border-radius: 6px; + color: var(--text); + padding: 6px 8px; + font-family: inherit; + font-size: 13px; +} + +input[type="text"]:focus, +input[type="number"]:focus { + outline: none; + border-color: var(--accent); +} + +.btn-row { display: flex; - flex-direction: column; - background: #fff; + gap: 6px; + margin-bottom: 8px; +} + +.btn-row button { + flex: 1; +} + +/* ---------- lists ---------- */ + +.item-list { + list-style: none; + margin: 0; + padding: 0; + overflow-y: auto; + min-height: 0; + max-height: 220px; +} + +.item-list.grow-list { + flex: 1 1 auto; + max-height: none; } -.top-bar { - padding: 10px 20px; - border-bottom: 1px solid #ccc; - background: #fafafa; + +.item-list li { display: flex; - gap: 10px; align-items: center; - flex-wrap: wrap; - flex-shrink: 0; + gap: 8px; + padding: 5px 8px; + border-radius: 6px; + cursor: pointer; + user-select: none; + border: 1px solid transparent; } -.image-container { - flex: 1; - position: relative; - background: #eee; - overflow: auto; - cursor: crosshair; + +.item-list li:hover { + background: rgba(255, 255, 255, 0.05); } -.bounding-box { - position: absolute; - pointer-events: none; +.item-list li.active { + background: rgba(79, 140, 255, 0.14); + border-color: var(--accent); } -.label-tag { - position: absolute; - top: -20px; - left: 0; - color: white; - padding: 2px 6px; - font-size: 12px; - font-weight: bold; + +.color-dot { + width: 11px; + height: 11px; border-radius: 3px; - pointer-events: none; + flex: 0 0 auto; } -.label-item { - background: #e0e0e0; - padding: 5px 10px; - margin-bottom: 5px; - border-radius: 4px; - cursor: pointer; + +.item-name { + flex: 1; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + +.item-badge { + color: var(--text-dim); + font-size: 11px; + font-variant-numeric: tabular-nums; } -.label-item.active { - outline: 2px solid #2196f3; + +.class-id { + color: var(--text-dim); + font-size: 10px; + font-family: var(--mono); + flex: 0 0 auto; } -.upload-btn, -.nav-btn { - margin-top: 5px; - padding: 5px 10px; - font-size: 14px; - cursor: pointer; + +.thumb { + width: 34px; + height: 24px; + object-fit: cover; + border-radius: 4px; + background: var(--bg); + flex: 0 0 auto; } -input[type="text"] { - width: 100%; - padding: 5px; + +.del-btn { + padding: 0 4px; + color: var(--text-dim); font-size: 14px; - box-sizing: border-box; + line-height: 1; + visibility: hidden; } -input[type="range"] { - width: 300px; +.item-list li:hover .del-btn { + visibility: visible; } -.crosshair-line { - position: absolute; - pointer-events: none; - z-index: 10; -} - -#verticalLine { - width: 3px; - height: 300%; - top: 0; - background-image: repeating-linear-gradient( - to bottom, - rgba(255, 0, 0, 0.6) 0px, - rgba(255, 0, 0, 0.6) 5px, - transparent 5px, - transparent 10px - ); -} - -#horizontalLine { - height: 3px; - width: 300%; - left: 0; - background-image: repeating-linear-gradient( - to right, - rgba(255, 0, 0, 0.6) 0px, - rgba(255, 0, 0, 0.6) 5px, - transparent 5px, - transparent 10px - ); -} -#imageWrapper { - position: relative; - display: inline-block; - max-width: none; - max-height: none; - transform-origin: top left; + +.del-btn:hover { + color: var(--danger); } -#imageDisplay { - display: block; - max-width: none; - height: auto; - pointer-events: none; - user-drag: none; +.muted { + color: var(--text-dim); + font-size: 12px; +} + +.mono { + font-family: var(--mono); + font-size: 11px; +} + +/* ---------- stage ---------- */ + +#stage { + position: relative; + flex: 1 1 auto; + min-width: 0; + background: var(--bg); + background-image: radial-gradient(var(--border) 1px, transparent 1px); + background-size: 22px 22px; } -#overlay { +#canvas { position: absolute; - top: 0; - left: 0; + inset: 0; width: 100%; height: 100%; + touch-action: none; +} + +#emptyState { + position: absolute; + inset: 0; + display: flex; + flex-direction: column; + align-items: center; + justify-content: center; + gap: 2px; + color: var(--text-dim); pointer-events: none; } -.bounding-box { +#dropHint { position: absolute; - pointer-events: auto; - box-sizing: border-box; + inset: 12px; + display: flex; + align-items: center; + justify-content: center; + border: 2px dashed var(--accent); + border-radius: 12px; + background: rgba(79, 140, 255, 0.08); + color: var(--accent); + font-size: 16px; + font-weight: 600; + pointer-events: none; + z-index: 5; +} + +.hidden { + display: none !important; +} + +/* ---------- progress ---------- */ + +#progress { + position: absolute; + left: 50%; + bottom: 18px; + transform: translateX(-50%); + width: 340px; + max-width: 80%; + background: var(--bg-raised); + border: 1px solid var(--border); + border-radius: 10px; + padding: 10px 14px 12px; + box-shadow: 0 6px 24px rgba(0, 0, 0, 0.45); + z-index: 20; +} + +.progress-header { + display: flex; + justify-content: space-between; + gap: 10px; + margin-bottom: 7px; + font-size: 12px; +} + +#progressPct { + color: var(--text-dim); + font-variant-numeric: tabular-nums; +} + +.progress-track { + position: relative; + height: 6px; + border-radius: 3px; + background: var(--bg); + overflow: hidden; +} + +#progressFill { + position: absolute; + inset: 0 auto 0 0; + width: 0%; + border-radius: 3px; + background: var(--accent); + transition: width 0.15s ease-out; +} + +#progressFill.indeterminate { + width: 35%; + animation: progress-slide 1.1s ease-in-out infinite; +} + +@keyframes progress-slide { + 0% { + left: -35%; + } + + 100% { + left: 100%; + } +} + +/* ---------- status bar ---------- */ + +#statusbar { + display: flex; + align-items: center; + gap: 16px; + padding: 5px 14px; + background: var(--bg-panel); + border-top: 1px solid var(--border); + color: var(--text-dim); + flex: 0 0 auto; + white-space: nowrap; + overflow: hidden; +} + +#saveStatus.ok { + color: var(--ok); +} + +#saveStatus.busy { + color: var(--text-dim); +} + +#saveStatus.error { + color: var(--danger); +} + +/* ---------- dialogs ---------- */ + +dialog { + background: var(--bg-panel); + color: var(--text); + border: 1px solid var(--border); + border-radius: 12px; + padding: 20px; + width: 380px; + max-width: 90vw; +} + +dialog::backdrop { + background: rgba(0, 0, 0, 0.55); +} + +dialog h3 { + margin: 0 0 12px; + font-size: 15px; +} + +dialog form { + display: flex; + flex-direction: column; + gap: 10px; +} + +dialog label { + display: flex; + flex-direction: column; + gap: 4px; + color: var(--text-dim); + font-size: 12px; + flex: 1; +} + +dialog label.check { + flex-direction: row; + align-items: center; + gap: 8px; + color: var(--text); + font-size: 13px; +} + +.field-row { + display: flex; + gap: 10px; +} + +input[type="color"] { + border: 1px solid var(--border); + border-radius: 6px; + background: var(--bg); + height: 30px; + padding: 2px; +} + +.dialog-actions { + display: flex; + justify-content: flex-end; + gap: 8px; + margin-top: 6px; +} + +/* ---------- toasts ---------- */ + +#toasts { + position: fixed; + bottom: 42px; + right: 16px; + display: flex; + flex-direction: column; + gap: 8px; + z-index: 100; + max-width: 420px; +} + +.toast { + background: var(--bg-raised); + border: 1px solid var(--border); + border-left: 3px solid var(--accent); + border-radius: 8px; + padding: 10px 14px; + box-shadow: 0 6px 24px rgba(0, 0, 0, 0.4); + animation: toast-in 0.18s ease-out; + overflow-wrap: anywhere; +} + +.toast.error { + border-left-color: var(--danger); +} + +.toast.ok { + border-left-color: var(--ok); +} + +.toast .mono { + color: var(--text-dim); + display: block; + margin-top: 3px; +} + +@keyframes toast-in { + from { + transform: translateY(8px); + opacity: 0; + } + + to { + transform: translateY(0); + opacity: 1; + } } diff --git a/src/dynsight/_internal/vision/vision.py b/src/dynsight/_internal/vision/vision.py index 22510407..dbaed00a 100644 --- a/src/dynsight/_internal/vision/vision.py +++ b/src/dynsight/_internal/vision/vision.py @@ -4,7 +4,7 @@ import logging from pathlib import Path -from typing import TYPE_CHECKING, Callable +from typing import TYPE_CHECKING, Callable, cast import numpy as np import torch @@ -209,22 +209,27 @@ def predict( The maximum number of detections for a single frame / image. """ - self.prediction_results = self.model.predict( - source=self.source, - save=True, - save_txt=False, - save_conf=True, - show_labels=show_labels, - name=prediction_title, - project=self.output_path, - device=self.device, - augment=augment, - agnostic_nms=agnostic_nms, - classes=class_filter, - conf=confidence, - iou=iou, - imgsz=imgsz, - max_det=max_det, + # Without stream=True, ultralytics always returns a list of + # Results; the cast narrows the wider annotated return union. + self.prediction_results = cast( + "list[Results]", + self.model.predict( + source=self.source, + save=True, + save_txt=False, + save_conf=True, + show_labels=show_labels, + name=prediction_title, + project=self.output_path, + device=self.device, + augment=augment, + agnostic_nms=agnostic_nms, + classes=class_filter, + conf=confidence, + iou=iou, + imgsz=imgsz, + max_det=max_det, + ), ) def create_dataset_from_predictions( diff --git a/tests/test_regressions.py b/tests/test_regressions.py new file mode 100644 index 00000000..772c74f0 --- /dev/null +++ b/tests/test_regressions.py @@ -0,0 +1,147 @@ +"""Regression tests for three bugs found while writing the tutorials. + +1. Descriptors built on a neighbor list ignored the ``Trj`` slice and + walked off the end of it (``IndexError``). +2. ``spatialaverage`` indexed the descriptor with the trajectory's frame + count, so a descriptor defined on frame pairs (LENS, timeSOAP) raised a + bare ``IndexError`` from inside a worker process. +3. ``track_xyz`` returned a ``Trj`` built on a file whose frames may hold + different numbers of objects, which raises ``EOFError`` as soon as any + descriptor is computed on it. Planar tracked data also made + ``compute_lens`` infer a zero-thickness box and divide by zero. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +from dynsight.analysis import spatialaverage +from dynsight.track import track_xyz +from dynsight.trajectory import Trj + +SYSTEMS = Path(__file__).resolve().parent / "systems" +_LJ_N_PARTICLES = 5 +TRJ_2D = ( + Path(__file__).resolve().parent.parent + / "docs/source/_static/ex_test_files/trajectory.xyz" +) + + +def _trj_2d() -> Trj: + return Trj.init_from_xyz(traj_file=TRJ_2D, dt=1.0) + + +def test_orientational_op_on_a_sliced_trj() -> None: + """The descriptor must follow the slice, not the whole trajectory.""" + trj = _trj_2d() + _, psi_full = trj.get_orientational_op(r_cut=3.0, order=6) + + n_frames = 10 + sliced = trj.with_slice(slice(0, n_frames, 1)) + neigcounts, _ = sliced.get_coord_number(r_cut=3.0) + _, psi = sliced.get_orientational_op( + r_cut=3.0, order=6, neigcounts=neigcounts + ) + + assert psi.dataset.shape == (trj.n_atoms, n_frames) + assert np.allclose(psi.dataset, psi_full.dataset[:, :n_frames]) + + +def test_velocity_alignment_on_a_sliced_trj() -> None: + trj = _trj_2d() + n_frames = 10 + sliced = trj.with_slice(slice(0, n_frames, 1)) + neigcounts, _ = sliced.get_coord_number(r_cut=3.0) + _, phi = sliced.get_velocity_alignment(r_cut=3.0, neigcounts=neigcounts) + + # No velocities in an .xyz: displacements are used, hence n_frames - 1. + assert phi.dataset.shape == (trj.n_atoms, n_frames - 1) + + +def test_mismatched_neighbour_list_raises_clearly() -> None: + """A neighbor list from another slice must fail loudly, not silently.""" + trj = _trj_2d() + neigcounts, _ = trj.get_coord_number(r_cut=3.0) + sliced = trj.with_slice(slice(0, 10, 1)) + + with pytest.raises(ValueError, match="neigh_list_per_frame covers"): + sliced.get_orientational_op(r_cut=3.0, order=6, neigcounts=neigcounts) + + +def test_spatial_average_frame_mismatch_raises_clearly() -> None: + """LENS is one frame shorter than its trajectory: say so, don't crash.""" + trj = _trj_2d() + descriptor = np.zeros((trj.n_atoms, trj.n_frames - 1)) + + with pytest.raises(ValueError, match="descriptor_array covers"): + spatialaverage( + universe=trj.universe, + descriptor_array=descriptor, + selection="all", + r_cut=3.0, + ) + + +def test_lens_on_planar_data_without_a_box(tmp_path: Path) -> None: + """A flat system has zero extent along z; the box must still be usable. + + This is exactly the shape of the data ``dynsight.vision`` produces: + pixel coordinates in a plane, with no simulation box. + """ + rng = np.random.default_rng(42) + n_atoms, n_frames = 30, 6 + xy = rng.uniform(0.0, 40.0, size=(n_frames, n_atoms, 2)) + + planar = tmp_path / "planar.xyz" + with planar.open("w") as file: + for frame in range(n_frames): + file.write(f"{n_atoms}\nFrame {frame}\n") + for x, y in xy[frame]: + file.write(f"P {x:.5f} {y:.5f} 0.00000\n") + + trj = Trj.init_from_xyz(traj_file=planar, dt=1.0) + assert trj.universe.trajectory[0].dimensions is None + assert np.allclose(trj.get_coordinates("all")[:, :, 2], 0.0) + + lens = trj.get_lens(r_cut=10.0) + + assert lens.dataset.shape == (n_atoms, n_frames - 1) + assert np.all(np.isfinite(lens.dataset)) + + +def test_track_xyz_returns_a_trj_when_the_count_is_constant( + tmp_path: Path, +) -> None: + trj = track_xyz( + input_xyz=SYSTEMS / "lj_noid.xyz", + output_xyz=tmp_path / "tracked.xyz", + search_range=10, + ) + assert trj is not None + assert trj.n_atoms == _LJ_N_PARTICLES + + +def test_track_xyz_returns_none_on_a_ragged_file(tmp_path: Path) -> None: + """A variable object count cannot make a trajectory: return None.""" + ragged = tmp_path / "ragged.xyz" + ragged.write_text( + "3\nf0\n1.0 1.0 0.0\n5.0 1.0 0.0\n9.0 1.0 0.0\n" + "2\nf1\n1.2 1.0 0.0\n5.2 1.0 0.0\n" + "3\nf2\n1.4 1.0 0.0\n5.4 1.0 0.0\n9.4 1.0 0.0\n" + ) + output = tmp_path / "tracked.xyz" + + assert ( + track_xyz( + input_xyz=ragged, + output_xyz=output, + search_range=3, + memory=0, + ) + is None + ) + # The file is still written: it is a faithful record of the detections. + assert output.exists() diff --git a/tests/track/test_track.py b/tests/track/test_track.py index fcb576d2..ebf4ed8f 100644 --- a/tests/track/test_track.py +++ b/tests/track/test_track.py @@ -1,12 +1,17 @@ from __future__ import annotations from pathlib import Path +from typing import TYPE_CHECKING import numpy as np +import pytest from dynsight.track import track_xyz from dynsight.utilities import read_xyz +if TYPE_CHECKING: + from dynsight._internal.utilities.utilities import Col + def test_track_xyz(tmp_path: Path) -> None: original_dir = Path(__file__).resolve().parent @@ -24,3 +29,53 @@ def test_track_xyz(tmp_path: Path) -> None: ).to_numpy() assert arr1.shape == arr2.shape assert np.array_equal(arr1, arr2) + + +NAMED_COLS = 4 + + +def strip_names(input_xyz: Path, output_xyz: Path) -> None: + """Write a copy of an .xyz file without its name column.""" + lines = [] + for line in input_xyz.read_text().splitlines(): + parts = line.split() + if len(parts) == NAMED_COLS: + parts = parts[1:] + lines.append(" ".join(parts)) + output_xyz.write_text("\n".join(lines) + "\n") + + +def test_track_xyz_without_names(tmp_path: Path) -> None: + # The name column is optional in the input file. + original_dir = Path(__file__).resolve().parent + file_with_id = original_dir / "../systems/lj_id.xyz" + + nameless = tmp_path / "nameless.xyz" + strip_names(original_dir / "../systems/lj_noid.xyz", nameless) + + output = tmp_path / "trajectory.xyz" + track_xyz(input_xyz=nameless, output_xyz=output, search_range=10) + + cols_order: list[Col] = ["name", "x", "y", "z", "ID"] + tracked = read_xyz(input_xyz=output, cols_order=cols_order) + expected = read_xyz(input_xyz=file_with_id, cols_order=cols_order) + + # Positions and IDs match the run on the file with names, and the + # output is a valid .xyz: the missing names get a placeholder. + compared = ["frame", "x", "y", "z", "ID"] + assert tracked.shape == expected.shape + assert np.array_equal( + tracked[compared].to_numpy(), expected[compared].to_numpy() + ) + assert set(tracked["name"]) == {"C"} + + +def test_track_xyz_invalid_format(tmp_path: Path) -> None: + invalid = tmp_path / "invalid.xyz" + invalid.write_text("2\ncomment\n1.0 2.0\n3.0 4.0\n") + with pytest.raises(ValueError, match=r"Error in the \.xyz format"): + track_xyz( + input_xyz=invalid, + output_xyz=tmp_path / "out.xyz", + search_range=10, + ) diff --git a/tests/vision/test_label_tool.py b/tests/vision/test_label_tool.py new file mode 100644 index 00000000..a1dc0d49 --- /dev/null +++ b/tests/vision/test_label_tool.py @@ -0,0 +1,272 @@ +from __future__ import annotations + +import io +import json +import threading +import urllib.error +import urllib.request +from typing import TYPE_CHECKING, Any + +import pytest +import yaml +from PIL import Image + +from dynsight._internal.vision.label_tool import ( + _LabelToolServer, + _safe_name, + _split_count, + _Workspace, + export_dataset, + load_session_file, + save_session_file, + synthesize_dataset, +) + +if TYPE_CHECKING: + from pathlib import Path + + +def make_image_bytes( + width: int = 64, height: int = 48, color: str = "red" +) -> bytes: + buffer = io.BytesIO() + Image.new("RGB", (width, height), color).save(buffer, format="PNG") + return buffer.getvalue() + + +@pytest.fixture +def workspace(tmp_path: Path) -> _Workspace: + ws = _Workspace(tmp_path / "ws") + for i in range(4): + ws.add_image(f"img_{i}.png", make_image_bytes()) + return ws + + +def make_session() -> dict[str, Any]: + return { + "labels": [ + {"name": "particle", "color": "#ff0000"}, + {"name": "aggregate", "color": "#00ff00"}, + ], + "annotations": { + "img_0.png": [ + {"label": "aggregate", "x": 4, "y": 6, "w": 20, "h": 10}, + {"label": "particle", "x": 30, "y": 20, "w": 10, "h": 12}, + ], + "img_1.png": [ + {"label": "particle", "x": 0, "y": 0, "w": 8, "h": 8}, + ], + }, + } + + +def test_safe_name_blocks_traversal() -> None: + assert _safe_name("../../etc/secret.png") == "secret.png" + assert _safe_name("a b/c?.png") == "c_.png" + with pytest.raises(ValueError, match="Invalid file name"): + _safe_name("...") + + +def test_split_count_keeps_val_nonempty() -> None: + assert _split_count(10, 0.8) == 8 # noqa: PLR2004 + assert _split_count(2, 0.99) == 1 + assert _split_count(2, 0.01) == 1 + assert _split_count(1, 0.8) == 1 + + +def test_add_image_rejects_invalid_data(tmp_path: Path) -> None: + ws = _Workspace(tmp_path / "ws") + with pytest.raises(ValueError, match="not a readable image"): + ws.add_image("bad.png", b"not an image") + assert ws.list_images() == [] + with pytest.raises(ValueError, match="Unsupported image format"): + ws.add_image("file.txt", b"hello") + + +def test_session_file_roundtrip(tmp_path: Path) -> None: + session = make_session() + path = save_session_file(session, tmp_path / "sub" / "session.json") + assert path.is_file() + assert load_session_file(path) == session + # A directory gets a default file name, missing suffixes are added. + assert save_session_file(session, tmp_path).name == "session.json" + assert save_session_file(session, tmp_path / "named").name == ( + "named.json" + ) + with pytest.raises(ValueError, match="file path is required"): + save_session_file(session, "") + with pytest.raises(ValueError, match="not found"): + load_session_file(tmp_path / "missing.json") + + +def test_export_dataset_layout(workspace: _Workspace) -> None: + result = export_dataset( + workspace, + make_session(), + name="my_dataset", + train_split=0.75, + shuffle=True, + seed=42, + ) + dataset = workspace.root / "my_dataset" + assert result["num_train"] == 3 # noqa: PLR2004 + assert result["num_val"] == 1 + + with (dataset / "dataset.yaml").open() as f: + content = yaml.safe_load(f) + assert content["path"] == str(dataset.resolve()) + assert content["train"] == "images/train" + assert content["val"] == "images/val" + assert content["nc"] == 2 # noqa: PLR2004 + assert content["names"] == ["particle", "aggregate"] + + images = sorted(p.name for p in dataset.glob("images/*/*")) + labels = sorted(p.name for p in dataset.glob("labels/*/*")) + assert images == [f"img_{i}.png" for i in range(4)] + assert labels == [f"img_{i}.txt" for i in range(4)] + + # Every image has a label file in the matching split folder. + for img_path in dataset.glob("images/*/*"): + split = img_path.parent.name + lbl = dataset / "labels" / split / (img_path.stem + ".txt") + assert lbl.is_file() + + +def test_export_dataset_stable_class_ids(workspace: _Workspace) -> None: + # "aggregate" is the second label: its class ID must be 1 even if + # it is the first annotation encountered. + export_dataset(workspace, make_session(), name="ds", shuffle=False, seed=0) + dataset = workspace.root / "ds" + lines = ( + (dataset / "labels" / "train" / "img_0.txt") + .read_text() + .strip() + .splitlines() + ) + class_ids = [line.split()[0] for line in lines] + assert class_ids == ["1", "0"] + # YOLO boxes are normalized cx cy w h in [0, 1]. + for line in lines: + values = [float(v) for v in line.split()[1:]] + assert all(0.0 <= v <= 1.0 for v in values) + + +def test_export_dataset_errors(workspace: _Workspace) -> None: + with pytest.raises(ValueError, match="No labels defined"): + export_dataset(workspace, {"labels": []}, name="ds") + with pytest.raises(ValueError, match="train_split"): + export_dataset(workspace, make_session(), name="ds", train_split=1.5) + + +def test_export_dataset_custom_output( + workspace: _Workspace, tmp_path: Path +) -> None: + out = tmp_path / "elsewhere" + result = export_dataset( + workspace, make_session(), name="ds", output_dir=out + ) + assert result["path"] == str((out / "ds").resolve()) + assert (out / "ds" / "dataset.yaml").is_file() + + +def test_synthesize_dataset(workspace: _Workspace) -> None: + result = synthesize_dataset( + workspace, + make_session(), + name="synt", + num_images=5, + width=128, + height=96, + per_image=3, + train_split=0.8, + scale_range=(0.5, 1.5), + seed=7, + ) + dataset = workspace.root / "synt" + assert result["num_train"] == 4 # noqa: PLR2004 + assert result["num_val"] == 1 + train_images = list(dataset.glob("images/train/*.jpg")) + assert len(train_images) == 4 # noqa: PLR2004 + with Image.open(train_images[0]) as img: + assert img.size == (128, 96) + for lbl in dataset.glob("labels/*/*.txt"): + for line in lbl.read_text().splitlines(): + parts = line.split() + assert parts[0] in {"0", "1"} + assert all(0.0 <= float(v) <= 1.0 for v in parts[1:]) + + +def test_synthesize_requires_annotations(tmp_path: Path) -> None: + ws = _Workspace(tmp_path / "ws") + ws.add_image("img.png", make_image_bytes()) + with pytest.raises(ValueError, match="No annotations"): + synthesize_dataset(ws, {"labels": [], "annotations": {}}, name="s") + + +def test_http_api_roundtrip(tmp_path: Path) -> None: + server = _LabelToolServer(0, _Workspace(tmp_path / "ws")) + port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + base = f"http://127.0.0.1:{port}" + + def request( + path: str, method: str = "GET", data: bytes | None = None + ) -> dict[str, Any]: + req = urllib.request.Request( # noqa: S310 + base + path, data=data, method=method + ) + with urllib.request.urlopen(req) as response: # noqa: S310 + return json.loads(response.read()) + + try: + info = request("/api/images?name=img.png", "POST", make_image_bytes()) + assert info == {"name": "img.png", "width": 64, "height": 48} + + # Edits are mirrored to the server memory, without disk writes. + session = make_session() + request("/api/sync", "POST", json.dumps(session).encode("utf-8")) + + state = request("/api/state") + assert [img["name"] for img in state["images"]] == ["img.png"] + assert state["labels"] == session["labels"] + assert state["dirty"] is True + assert state["session_path"] is None + + # No long operation running: the progress endpoint is idle. + assert request("/api/progress") == {"active": False} + + # Saving requires an explicit path. + with pytest.raises(urllib.error.HTTPError): + request("/api/session", "POST", b"{}") + session_file = tmp_path / "saved" / "session.json" + saved = request( + "/api/session", + "POST", + json.dumps({"path": str(session_file)}).encode("utf-8"), + ) + assert saved["path"] == str(session_file) + assert session_file.is_file() + assert request("/api/state")["dirty"] is False + + loaded = request( + "/api/session/load", + "POST", + json.dumps({"path": str(session_file)}).encode("utf-8"), + ) + assert loaded == session + + export = request( + "/api/export", + "POST", + json.dumps({"name": "ds", "seed": 1}).encode("utf-8"), + ) + assert (tmp_path / "ws" / "ds" / "dataset.yaml").is_file() + assert export["num_train"] == 1 + + request("/api/images?name=img.png", "DELETE") + assert request("/api/state")["images"] == [] + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5)