update docs for 3D

docs/source/conf.py

@@ -12,7 +12,7 @@
 project = 'Spotiflow'
 copyright = '2023, Albert Dominguez Mantes'
 author = 'Albert Dominguez Mantes'
-release = '0.1'
+release = '0.4.0'
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

docs/source/napari.rst

@@ -13,14 +13,14 @@ The plugin will then be available in the napari GUI under the name "Spotiflow wi
     :width: 700
     :align: center
 
-The plugin has two modes: for images (``2D``) and movies (``2D+t``), which can be toggled using the ``2D`` and ``2D+t`` buttons in the GUI. Other options can be tuned, `e.g.` rescaling the image.
+The plugin allows running on two modes: for images (``2D``) and volumes (``3D``), which can be toggled using the corresponding  buttons in the GUI. You can also run on movies by setting the appropriate axis order (should be leading with a `T`).
 
 Upon pressing the button ``Run``, The plugin will create a ``Points`` layer containing the predicted spots:
 
 .. image:: ./_static/spotiflow_napari_preds.png
     :width: 700
     :align: center
 
-If the option ``Show CNN output`` is checked, the plugin will also create an ``Image`` layer containing the heatmap output of the CNN.
+If the option ``Show CNN output`` is checked, the plugin will also create two ``Image`` layers containing the heatmap output of the CNN as well as the stereographic flow.
 
-Finally, the plugin includes one sample image per dataset. These samples can be loaded from the ``File`` menu (``File -> Open sample -> napari-spotiflow``). You can try the plugin with these samples to get a better idea of how it works!
+Finally, the plugin includes two sample 2D images (HybISS and Terra) as well as a synthetic 3D volume. These samples can be loaded from the ``File`` menu (``File -> Open sample -> napari-spotiflow``). You can try the plugin with these samples to get a better idea of how it works!

docs/source/train.rst

@@ -31,7 +31,16 @@ The actual naming of the files is not important, but the ``.csv`` and ``.tif`` f
     252.99, 307.97
     ...
 
-The column names can also be `axis-0` (instead of `y`) and `axis-1` instead of `x`.
+The column names can also be `axis-0` (instead of `y`) and `axis-1` instead of `x`. For the 3D case, the format is similar but with an additional column corresponding to the `z` coordinate:
+
+.. code-block::
+
+    z,y,x 
+    12.4,42.3,24.24
+    61.2,252.99, 307.97
+    ...
+
+In this case, you can also use `axis-0`, `axis-1`, and `axis-2` instead of `z`, `y`, and `x`, respectively.
 
 
 Basic training
@@ -65,6 +74,24 @@ You can then load it by simply calling:
 
     model = Spotiflow.from_folder("/my/trained/model")
 
+
+In the 3D case, you should initialize a :py:mod:`spotiflow.model.config.SpotiflowModelConfig` object and pass it to the `Spotiflow` constructor with the appropriate parameter set (see other options for the configuration at the end of the section):
+
+.. code-block:: python
+    # Same imports as before
+    from spotiflow.model import SpotiflowModelConfig
+
+    # Create the model config
+    model_config = SpotiflowModelConfig(
+        is_3d=True,
+        grid=2, # subsampling factor for prediction
+        # you can pass other arguments here
+    )
+
+    model = Spotiflow(model_config)
+    # Train and save the model as before
+
+
 Customizing the training
 ^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -75,6 +102,7 @@ You can also pass other parameters relevant for training to the `fit` method. Fo
     train_config = {
         "num_epochs": 100,
         "learning_rate": 0.001,
+        "smart_crop": True,
         # other parameters
     }
 
@@ -89,7 +117,7 @@ You can also pass other parameters relevant for training to the `fit` method. Fo
     )
 
 
-In order to change the model architecture (`e.g.` number of input/output channels, number of layers, variance for the heatmap generation, etc.), you can create a :py:mod:`spotiflow.model.config.SpotiflowModelConfig` object and populate it accordingly. Then you can pass it to the `Spotiflow` constructor. For example, if our image is RGB and we need the network to use 3 input channels, we can do the following:
+In order to change the model architecture (`e.g.` number of input/output channels, number of layers, variance for the heatmap generation, etc.), you can create a :py:mod:`spotiflow.model.config.SpotiflowModelConfig` object and populate it accordingly. Then you can pass it to the `Spotiflow` constructor (note that this is necessary for 3D). For example, if our image is RGB and we need the network to use 3 input channels, we can do the following:
 
 .. code-block:: python