Vision revision and track

.github/workflows/tests.yaml

@@ -60,7 +60,7 @@ jobs:
               with:
                   python-version: ${{ matrix.python-version }}
                   cache: "pip"
-            - run: pip install -e '.[dev]'
+            - run: pip install -e '.[dev, track]'
             - run: pytest --cov=src --cov-report term-missing
     doctest:
         runs-on: ubuntu-latest
@@ -74,5 +74,5 @@ jobs:
               with:
                   python-version: ${{ matrix.python-version }}
                   cache: "pip"
-            - run: pip install -e '.[dev]'
-            - run: make -C docs doctest
+            - run: pip install -e '.[dev, track]'
+            - run: make -C docs doctest

.readthedocs.yaml

@@ -1,16 +1,17 @@
 version: 2
 
 build:
-    os: ubuntu-22.04
-    tools:
-      python: "3.10"
+  os: ubuntu-22.04
+  tools:
+    python: "3.10"
 
 sphinx:
-    configuration: docs/source/conf.py
+  configuration: docs/source/conf.py
 
 python:
-    install:
-        - method: pip
-          path: .
-          extra_requirements:
-            - dev
+  install:
+    - method: pip
+      path: .
+      extra_requirements:
+        - dev
+        - track

docs/source/_templates/class.rst

@@ -18,14 +18,3 @@
    {%- endfor %}
    {% endif %}
    {% endblock %}
-
-   {% block attributes %}
-   {% if attributes %}
-   .. rubric:: {{ _('Attributes') }}
-
-   .. autosummary::
-   {% for item in attributes %}
-      ~{{ name }}.{{ item }}
-   {%- endfor %}
-   {% endif %}
-   {% endblock %}

docs/source/index.rst

@@ -9,6 +9,7 @@
    tICA <tica>
    onion clustering <onion>
    vision <vision>
+   track <track>
    analysis <analysis>
    data processing <data_processing>
    HDF5er <hdf5er>
@@ -74,7 +75,7 @@ or with conda::
 If you want to use the :mod:`dynsight.vision` and :mod:`dynsight.track` modules
 you will need to install a series of packages. This can be done with with pip::
 
-  $ pip install ultralytics PyYAML
+  $ pip install ultralytics PyYAML pandas trackpy
 
 
 How to contribute

docs/source/track.rst

@@ -0,0 +1,14 @@
+track
+=====
+
+This module tracks particle trajectories in `.xyz` files that lack explicit
+particle IDs, assigning consistent identities across frames.
+It is especially useful for outputs generated by the :doc:`vision` module.
+
+Functions
+---------
+
+.. toctree::
+  :maxdepth: 1
+
+  track_xyz <_autosummary/dynsight.track.track_xyz>

docs/source/vision.rst

@@ -1,7 +1,43 @@
 Vision
 ======
 
-Code for computer vision and trajectory extraction from experimental videos.
+Code that contains useful tools to obtain trajectories from video files.
+The extraction is possible thanks to the optimization of a Convolutional Neural
+Network (CNN) model training to
+detect specific objects in the video. Once the positions are recovered, a
+tracking algorithm (:doc:`track`) links each object with its own identity number.
+
+.. figure:: _static/vision/vision_scheme.jpeg
+   :alt: Schematic representation of a typical dynsight.vision application.
+   :align: center
+
+   Schematic representation of a typical `dynsight.vision` application.
+
+The object detection is managed by the Ultralytics_ external library. If you
+use this tool, please cite them by referring to this link_.
+
+.. _Ultralytics: https://www.ultralytics.com
+.. _link: https://docs.ultralytics.com/models/yolo11/#citations-and-acknowledgements
+
+--------------
+The Vision GUI
+--------------
+
+The `dynsight.vision` module includes an internal Graphical User Interface (GUI)
+designed to assist users in preparing the training items required for object
+detection. It only requires an *initial guess* of the objects present in
+the video, which is then used to generate a synthetic dataset to bootstrap the
+training process.
+
+.. tip::
+
+   Applying some basic color correction to the input video (especially converting
+   it to grayscale with enhanced contrast) can significantly improve the quality
+   of the initial model.
+
+.. image:: _static/vision/vision_gui.gif
+   :alt: Usage of the GUI.
+   :align: center
 
 -----
 Usage
@@ -12,3 +48,6 @@ Usage
 
   Video <_autosummary/dynsight.vision.Video.rst>
   Detect <_autosummary/dynsight.vision.Detect.rst>
+
+Once obtained the particle positions, is possible to obtain the trajectories of
+each particle by using the :doc:`track` module.

examples/video_detection.py

@@ -113,10 +113,25 @@ def main() -> None:
         )
     if args.predict:
         detection.predict_frames(model_path=args.detect_model)
+        xyz_trajectory = Path.cwd() / "trajectory.xyz"
         detection.compute_xyz(
             prediction_folder_path=Path("prediction"),
             output_path=Path.cwd(),
         )
+        tr_xyz_trajectory = Path.cwd() / "tracked_trajectory.xyz"
+        trj = dynsight.track.track_xyz(
+            input_xyz=xyz_trajectory,
+            output_xyz=tr_xyz_trajectory,
+        )
+        # Usage example of the trj object after tracking.
+        lens_descriptor = trj.get_lens(r_cut=5)
+        lens_descriptor.dump_to_json(output_project)
+        # Compute clustering
+        onion_lens = lens_descriptor.get_onion(delta_t=2)
+        onion_lens.plot_output(
+            file_path=args.output,
+            data_insight=lens_descriptor,
+        )
 
 
 if __name__ == "__main__":

pyproject.toml

@@ -30,6 +30,7 @@ dev = [
     "furo",
     "numba",
 ]
+track = ["trackpy", "pandas"]
 
 [project.urls]
 github = "https://github.com/GMPavanLab/dynsight"
@@ -110,6 +111,8 @@ module = [
     'Pillow.*',
     'ultralytics.*',
     'opencv-python.*',
+    'pandas.*',
+    'trackpy.*',
     'cv2.*',
 ]
 ignore_missing_imports = true

src/dynsight/__init__.py

@@ -8,6 +8,7 @@
     onion,
     soap,
     tica,
+    track,
     trajectory,
     utilities,
     vision,
@@ -21,6 +22,7 @@
     "onion",
     "soap",
     "tica",
+    "track",
     "trajectory",
     "utilities",
     "vision",

src/dynsight/_internal/track/track.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import pandas as pd
+import trackpy as tp
+
+from dynsight.trajectory import Trj
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s | %(levelname)s | %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+
+def track_xyz(
+    input_xyz: Path,
+    output_xyz: Path,
+    search_range: float = 10,
+    memory: int = 1,
+    adaptive_stop: None | float = 0.95,
+    adaptive_step: None | float = 0.5,
+) -> Trj:
+    """Track particles from an `.xyz` file and write a new file with IDs.
+
+    The input `.xyz` is assumed to contain only raw 3D coordinates
+    (without atom labels/identity), and each frame begins with a line
+    indicating the number of atoms, followed by a comment line, then a list of
+    positions. Each frame in the input file must follow this structure::
+
+        <number of atoms>
+        comment line
+        <x> <y> <z>
+        <x> <y> <z>
+        ...
+        <x> <y> <z>
+
+    The output file will have the same structure, but each line will start
+    with the tracked particle ID.
+
+    Parameters:
+        input_xyz:
+            Path to the input .xyz file containing positions only.
+
+        output_xyz:
+            Path where the output .xyz file with particle IDs will be saved.
+
+        search_range:
+            The maximum allowable displacement of objects between frames for
+            them to be considered the same particle.
+
+        memory:
+            The maximum number of frames during which an object can vanish,
+            then reppear nearby, and be considered the same particle.
+
+        adaptive_stop:
+            If not `None`, when encountering a region with too many candidate
+            links (subnet), retry by progressively reducing `search_range`
+            until the subnet is solvable. If `search_range` becomes less or
+            equal than the `adaptive_stop`, give up and raise a
+            `SubnetOversizeException`.
+
+        adaptive_step:
+            Factor by which the `search_range` is multiplied to reduce it
+            during adaptive search. Effective only if `adaptive_stop` is not
+            `None`.
+    """
+    if adaptive_stop is None and adaptive_step is not None:
+        msg = "adaptive_step is set but adaptive_stop is None."
+        raise ValueError(msg)
+    if adaptive_stop is not None and adaptive_step is None:
+        msg = "adaptive_stop is set but adaptive_step is None."
+        raise ValueError(msg)
+
+    input_xyz = Path(input_xyz)
+    output_xyz = Path(output_xyz)
+
+    if not input_xyz.exists():
+        msg = f"Input file not found: {input_xyz}"
+        raise FileNotFoundError(msg)
+
+    positions = _collect_positions(input_xyz=input_xyz)
+
+    if not {"frame", "x", "y", "z"}.issubset(positions.columns):
+        msg = "Error in the .xyz format. Each line must be <x> <y> <z>."
+        raise ValueError(msg)
+
+    linked = tp.link_df(
+        positions,
+        search_range=search_range,
+        memory=memory,
+        adaptive_step=adaptive_step,
+        adaptive_stop=adaptive_stop,
+    )
+
+    with output_xyz.open("w") as f:
+        for frame_num in sorted(linked["frame"].unique()):
+            frame_data = linked[linked["frame"] == frame_num]
+            f.write(f"{len(frame_data)}\n")
+            f.write(f"Frame {frame_num}\n")
+            for _, row in frame_data.iterrows():
+                pid = int(row["particle"])
+                x, y, z = row["x"], row["y"], row["z"]
+                f.write(f"{pid} {x:.6f} {y:.6f} {z:.6f}\n")
+
+    logger.info(f"Linked .xyz file written to: {output_xyz}")
+    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()
+
+    data = []
+    frame = -1
+    row = 0
+    dimensions = 3
+    for _ in range(len(lines)):
+        if row >= 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})
+            row += num_atoms
+        else:
+            row += 1
+
+    return pd.DataFrame(data)