diff --git a/src/napari_trackastra/__init__.py b/src/napari_trackastra/__init__.py
index f510a10..23e6fed 100644
--- a/src/napari_trackastra/__init__.py
+++ b/src/napari_trackastra/__init__.py
@@ -1,11 +1,4 @@
__version__ = "0.0.1"
-from ._sample_data import make_sample_data
-from ._widget import ExampleQWidget, ImageThreshold, threshold_autogenerate_widget, threshold_magic_widget
+from ._widget import Tracker
-__all__ = (
- "make_sample_data",
- "ExampleQWidget",
- "ImageThreshold",
- "threshold_autogenerate_widget",
- "threshold_magic_widget",
-)
+
diff --git a/src/napari_trackastra/_sample_data.py b/src/napari_trackastra/_sample_data.py
index 453db51..241618a 100644
--- a/src/napari_trackastra/_sample_data.py
+++ b/src/napari_trackastra/_sample_data.py
@@ -8,14 +8,17 @@
"""
from __future__ import annotations
+from pathlib import Path
import numpy
+import tifffile
+from trackastra import data
-def make_sample_data():
- """Generates an image"""
- # Return list of tuples
- # [(data1, add_image_kwargs1), (data2, add_image_kwargs2)]
- # Check the documentation for more information about the
- # add_image_kwargs
- # https://napari.org/stable/api/napari.Viewer.html#napari.Viewer.add_image
- return [(numpy.random.rand(512, 512), {})]
+def test_data_bacteria() -> list[tuple[numpy.ndarray, dict, str]]:
+ imgs, masks = data.test_data_bacteria()
+ return [(imgs, dict(name='img'), 'image'), (masks, dict(name='mask'), 'labels')]
+
+
+def test_data_hela() -> list[tuple[numpy.ndarray, dict, str]]:
+ imgs, masks = data.test_data_hela()
+ return [(imgs, dict(name='img'), 'image'), (masks, dict(name='mask'), 'labels')]
diff --git a/src/napari_trackastra/_tests/test_sample_data.py b/src/napari_trackastra/_tests/test_sample_data.py
deleted file mode 100644
index 71b37bf..0000000
--- a/src/napari_trackastra/_tests/test_sample_data.py
+++ /dev/null
@@ -1,7 +0,0 @@
-# from napari_trackastra import make_sample_data
-
-# add your tests here...
-
-
-def test_something():
- pass
diff --git a/src/napari_trackastra/_tests/test_widget.py b/src/napari_trackastra/_tests/test_widget.py
index 8b45194..95f0309 100644
--- a/src/napari_trackastra/_tests/test_widget.py
+++ b/src/napari_trackastra/_tests/test_widget.py
@@ -1,66 +1,21 @@
import numpy as np
+import napari
+from napari_trackastra._widget import Tracker
+from trackastra.data import test_data_bacteria
-from napari_trackastra._widget import (
- ExampleQWidget,
- ImageThreshold,
- threshold_autogenerate_widget,
- threshold_magic_widget,
-)
+def test_widget():
+ viewer = napari.Viewer()
+ img, mask = test_data_bacteria()
+ viewer.add_image(img)
+ viewer.add_labels(mask)
-def test_threshold_autogenerate_widget():
- # because our "widget" is a pure function, we can call it and
- # test it independently of napari
- im_data = np.random.random((100, 100))
- thresholded = threshold_autogenerate_widget(im_data, 0.5)
- assert thresholded.shape == im_data.shape
- # etc.
+ viewer.window.add_dock_widget(Tracker(viewer))
+
-# make_napari_viewer is a pytest fixture that returns a napari viewer object
-# you don't need to import it, as long as napari is installed
-# in your testing environment
-def test_threshold_magic_widget(make_napari_viewer):
- viewer = make_napari_viewer()
- layer = viewer.add_image(np.random.random((100, 100)))
+if __name__ == "__main__":
+ test_widget()
- # our widget will be a MagicFactory or FunctionGui instance
- my_widget = threshold_magic_widget()
-
- # if we "call" this object, it'll execute our function
- thresholded = my_widget(viewer.layers[0], 0.5)
- assert thresholded.shape == layer.data.shape
- # etc.
-
-
-def test_image_threshold_widget(make_napari_viewer):
- viewer = make_napari_viewer()
- layer = viewer.add_image(np.random.random((100, 100)))
- my_widget = ImageThreshold(viewer)
-
- # because we saved our widgets as attributes of the container
- # we can set their values without having to "interact" with the viewer
- my_widget._image_layer_combo.value = layer
- my_widget._threshold_slider.value = 0.5
-
- # this allows us to run our functions directly and ensure
- # correct results
- my_widget._threshold_im()
- assert len(viewer.layers) == 2
-
-
-# capsys is a pytest fixture that captures stdout and stderr output streams
-def test_example_q_widget(make_napari_viewer, capsys):
- # make viewer and add an image layer using our fixture
- viewer = make_napari_viewer()
- viewer.add_image(np.random.random((100, 100)))
-
- # create our widget, passing in the viewer
- my_widget = ExampleQWidget(viewer)
-
- # call our widget method
- my_widget._on_click()
-
- # read captured output and check that it's as we expected
- captured = capsys.readouterr()
- assert captured.out == "napari has 1 layers\n"
+ napari.run()
+
\ No newline at end of file
diff --git a/src/napari_trackastra/_widget.py b/src/napari_trackastra/_widget.py
index ed8c358..99da322 100644
--- a/src/napari_trackastra/_widget.py
+++ b/src/napari_trackastra/_widget.py
@@ -1,128 +1,128 @@
-"""
-This module contains four napari widgets declared in
-different ways:
-
-- a pure Python function flagged with `autogenerate: true`
- in the plugin manifest. Type annotations are used by
- magicgui to generate widgets for each parameter. Best
- suited for simple processing tasks - usually taking
- in and/or returning a layer.
-- a `magic_factory` decorated function. The `magic_factory`
- decorator allows us to customize aspects of the resulting
- GUI, including the widgets associated with each parameter.
- Best used when you have a very simple processing task,
- but want some control over the autogenerated widgets. If you
- find yourself needing to define lots of nested functions to achieve
- your functionality, maybe look at the `Container` widget!
-- a `magicgui.widgets.Container` subclass. This provides lots
- of flexibility and customization options while still supporting
- `magicgui` widgets and convenience methods for creating widgets
- from type annotations. If you want to customize your widgets and
- connect callbacks, this is the best widget option for you.
-- a `QWidget` subclass. This provides maximal flexibility but requires
- full specification of widget layouts, callbacks, events, etc.
-
-References:
-- Widget specification: https://napari.org/stable/plugins/guides.html?#widgets
-- magicgui docs: https://pyapp-kit.github.io/magicgui/
-
-Replace code below according to your needs.
-"""
-from typing import TYPE_CHECKING
-
-from magicgui import magic_factory
-from magicgui.widgets import CheckBox, Container, create_widget
-from qtpy.QtWidgets import QHBoxLayout, QPushButton, QWidget
-from skimage.util import img_as_float
-
-if TYPE_CHECKING:
- import napari
-
-
-# Uses the `autogenerate: true` flag in the plugin manifest
-# to indicate it should be wrapped as a magicgui to autogenerate
-# a widget.
-def threshold_autogenerate_widget(
- img: "napari.types.ImageData",
- threshold: "float",
-) -> "napari.types.LabelsData":
- return img_as_float(img) > threshold
-
-
-# the magic_factory decorator lets us customize aspects of our widget
-# we specify a widget type for the threshold parameter
-# and use auto_call=True so the function is called whenever
-# the value of a parameter changes
-@magic_factory(
- threshold={"widget_type": "FloatSlider", "max": 1}, auto_call=True
-)
-def threshold_magic_widget(
- img_layer: "napari.layers.Image", threshold: "float"
-) -> "napari.types.LabelsData":
- return img_as_float(img_layer.data) > threshold
+
+import torch
+import numpy as np
+
+
+import napari
+from magicgui import magic_factory, magicgui
+from magicgui.widgets import CheckBox, Container, create_widget, PushButton, FileEdit, ComboBox, RadioButtons
+from pathlib import Path
+from typing import List
+from napari.utils import progress
+import trackastra
+from trackastra.utils import normalize
+from trackastra.model import Trackastra
+from trackastra.tracking import graph_to_ctc, graph_to_napari_tracks
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def _track_function(model, imgs, masks, **kwargs):
+ print("Normalizing...")
+ imgs = np.stack([normalize(x) for x in imgs])
+ print("Tracking...")
+ track_graph = model.track(imgs, masks, mode="greedy",
+ max_distance=128,
+ progbar_class=progress,
+ **kwargs) # or mode="ilp"
+ # Visualise in napari
+ df, masks_tracked = graph_to_ctc(track_graph,masks,outdir=None)
+ napari_tracks, napari_tracks_graph, _ = graph_to_napari_tracks(track_graph)
+ return track_graph, masks_tracked, napari_tracks
+
+
+# logo = Path(__file__).parent/"resources"/"trackastra_logo_small.png"
+
+# @magicgui(call_button="track",
+# label_head=dict(widget_type="Label", label=f'Trackastra
'),
+# model_path={"label": "Model Path", "mode": "d"},
+# persist=True)
+# def track(label_head, img_layer: napari.layers.Image, mask_layer:napari.layers.Labels, model_path:Path, distance_costs:bool=False) -> List[napari.types.LayerDataTuple]:
+# if model_path.exists():
+# model = Trackastra.from_folder(model_path, device=device)
+# else:
+# model = Trackastra.from_pretrained(model_path.name, device=device)
+# imgs = np.asarray(img_layer.data)
+# masks = np.asarray(mask_layer.data)
+# track_graph, masks_tracked, napari_tracks = _track_function(model, imgs, masks, use_distance=distance_costs)
+# mask_layer.visible = False
+# return [(napari_tracks, dict(name='tracks',tail_length=5), "tracks"), (masks_tracked, dict(name='masks_tracked', opacity=0.3), "labels")]
+
# if we want even more control over our widget, we can use
# magicgui `Container`
-class ImageThreshold(Container):
+class Tracker(Container):
def __init__(self, viewer: "napari.viewer.Viewer"):
super().__init__()
self._viewer = viewer
- # use create_widget to generate widgets from type annotations
- self._image_layer_combo = create_widget(
- label="Image", annotation="napari.layers.Image"
- )
- self._threshold_slider = create_widget(
- label="Threshold", annotation=float, widget_type="FloatSlider"
- )
- self._threshold_slider.min = 0
- self._threshold_slider.max = 1
- # use magicgui widgets directly
- self._invert_checkbox = CheckBox(text="Keep pixels below threshold")
+ self._label = create_widget(widget_type="Label", label=f'Trackastra
')
+ self._image_layer = create_widget(label="Images", annotation="napari.layers.Image")
+
+ self._out_mask, self._out_tracks = None, None
+
+ self._mask_layer = create_widget(label="Masks", annotation="napari.layers.Labels")
+ self._model_type = RadioButtons(label="Model Type", choices=["Pretrained", "Custom"], orientation="horizontal", value="Pretrained")
+ self._model_pretrained = ComboBox(label="Pretrained Model",
+ choices=tuple(trackastra.model.pretrained._MODELS.keys()), value="general_2d")
+ self._model_path = FileEdit(label="Model Path", mode="d")
+ self._model_path.hide()
+ self._run_button = PushButton(label="Track")
+
+
+ self._model_type.changed.connect(self._model_type_changed)
+ self._model_pretrained.changed.connect(self._update_model)
+ self._model_path.changed.connect(self._update_model)
+ self._run_button.changed.connect(self._run)
- # connect your own callbacks
- self._threshold_slider.changed.connect(self._threshold_im)
- self._invert_checkbox.changed.connect(self._threshold_im)
# append into/extend the container with your widgets
self.extend(
[
- self._image_layer_combo,
- self._threshold_slider,
- self._invert_checkbox,
+ self._label,
+ self._image_layer,
+ self._mask_layer,
+ self._model_type,
+ self._model_pretrained,
+ self._model_path,
+ self._run_button,
]
)
- def _threshold_im(self):
- image_layer = self._image_layer_combo.value
- if image_layer is None:
- return
-
- image = img_as_float(image_layer.data)
- name = image_layer.name + "_thresholded"
- threshold = self._threshold_slider.value
- if self._invert_checkbox.value:
- thresholded = image < threshold
+ def _model_type_changed(self, event):
+ if event == "Pretrained":
+ self._model_pretrained.show()
+ self._model_path.hide()
else:
- thresholded = image > threshold
- if name in self._viewer.layers:
- self._viewer.layers[name].data = thresholded
+ self._model_pretrained.hide()
+ self._model_path.show()
+
+ def _run(self, event=None):
+ self._update_model()
+
+ if self.model is None:
+ raise ValueError("Model not loaded")
+
+ imgs = np.asarray(self._image_layer.value.data)
+ masks = np.asarray(self._mask_layer.value.data)
+ track_graph, masks_tracked, napari_tracks = _track_function(self.model, imgs, masks)
+ self._mask_layer.value.visible = False
+
+
+ lays = tuple(lay for lay in self._viewer.layers if lay.name=="masks_tracked")
+ if len(lays) > 0:
+ lays[0].data = masks_tracked
else:
- self._viewer.add_labels(thresholded, name=name)
-
-
-class ExampleQWidget(QWidget):
- # your QWidget.__init__ can optionally request the napari viewer instance
- # use a type annotation of 'napari.viewer.Viewer' for any parameter
- def __init__(self, viewer: "napari.viewer.Viewer"):
- super().__init__()
- self.viewer = viewer
-
- btn = QPushButton("Click me!")
- btn.clicked.connect(self._on_click)
-
- self.setLayout(QHBoxLayout())
- self.layout().addWidget(btn)
-
- def _on_click(self):
- print("napari has", len(self.viewer.layers), "layers")
+ self._viewer.add_labels(masks_tracked, name="masks_tracked")
+
+ lays = tuple(lay for lay in self._viewer.layers if lay.name=="tracks")
+ if len(lays) > 0:
+ lays[0].data = napari_tracks
+ else:
+ self._viewer.add_tracks(napari_tracks, name="tracks")
+
+ def _update_model(self, event=None):
+ if self._model_type.value == "Pretrained":
+ self.model = Trackastra.from_pretrained(self._model_pretrained.value, device=device)
+ else:
+ self.model = Trackastra.from_folder(self._model_path.value, device=device)
\ No newline at end of file
diff --git a/src/napari_trackastra/napari.yaml b/src/napari_trackastra/napari.yaml
index ddc79b3..e0efc72 100644
--- a/src/napari_trackastra/napari.yaml
+++ b/src/napari_trackastra/napari.yaml
@@ -2,36 +2,25 @@ name: napari-trackastra
display_name: trackastra
# use 'hidden' to remove plugin from napari hub search results
visibility: public
-# see https://napari.org/stable/plugins/manifest.html for valid categories
-categories: ["Annotation", "Segmentation", "Acquisition"]
+categories: ["Segmentation"]
contributions:
commands:
- - id: napari-trackastra.make_sample_data
- python_name: napari_trackastra._sample_data:make_sample_data
- title: Load sample data from trackastra
- - id: napari-trackastra.make_container_widget
- python_name: napari_trackastra:ImageThreshold
- title: Make threshold Container widget
- - id: napari-trackastra.make_magic_widget
- python_name: napari_trackastra:threshold_magic_widget
- title: Make threshold magic widget
- - id: napari-trackastra.make_function_widget
- python_name: napari_trackastra:threshold_autogenerate_widget
- title: Make threshold function widget
- - id: napari-trackastra.make_qwidget
- python_name: napari_trackastra:ExampleQWidget
- title: Make example QWidget
+ - id: napari-trackastra.test_data_bacteria
+ python_name: napari_trackastra._sample_data:test_data_bacteria
+ title: Sample bacteria images and masks
+ - id: napari-trackastra.test_data_hela
+ python_name: napari_trackastra._sample_data:test_data_hela
+ title: Sample bacteria images and masks
+ - id: napari-trackastra.track
+ python_name: napari_trackastra:Tracker
+ title: Create Plugin
sample_data:
- - command: napari-trackastra.make_sample_data
- display_name: trackastra
+ - command: napari-trackastra.test_data_bacteria
+ display_name: bacteria
key: unique_id.1
+ - command: napari-trackastra.test_data_hela
+ display_name: hela
+ key: unique_id.2
widgets:
- - command: napari-trackastra.make_container_widget
- display_name: Container Threshold
- - command: napari-trackastra.make_magic_widget
- display_name: Magic Threshold
- - command: napari-trackastra.make_function_widget
- autogenerate: true
- display_name: Autogenerate Threshold
- - command: napari-trackastra.make_qwidget
- display_name: Example QWidget
+ - command: napari-trackastra.track
+ display_name: Trackastra tracking