Add sample data

Author: constantinpape

flamingo_tools/model_utils.py

@@ -1,5 +1,5 @@
 import os
-from typing import Optional, Union
+from typing import Dict, Optional, Union
 
 import pooch
 import torch
@@ -113,3 +113,49 @@ def get_model(model_type: str, device: Optional[Union[str, torch.device]] = None
     model = torch.load(model_path, weights_only=False)
     model.to(device)
     return model
+
+
+def get_default_tiling() -> Dict[str, Dict[str, int]]:
+    """Determine the tile shape and halo depending on the available VRAM.
+
+    Returns:
+        The default tiling settings for the available computational resources.
+    """
+    if torch.cuda.is_available():
+        # The default halo size.
+        halo = {"x": 64, "y": 64, "z": 16}
+
+        # Determine the GPU RAM and derive a suitable tiling.
+        vram = torch.cuda.get_device_properties(0).total_memory / 1e9
+
+        if vram >= 80:
+            tile = {"x": 640, "y": 640, "z": 80}
+        elif vram >= 40:
+            tile = {"x": 512, "y": 512, "z": 64}
+        elif vram >= 20:
+            tile = {"x": 352, "y": 352, "z": 48}
+        elif vram >= 10:
+            tile = {"x": 256, "y": 256, "z": 32}
+            halo = {"x": 64, "y": 64, "z": 8}  # Choose a smaller halo in z.
+        else:
+            raise NotImplementedError(f"Infererence with a GPU with {vram} GB VRAM is not supported.")
+
+        tiling = {"tile": tile, "halo": halo}
+        print(f"Determined tile size for CUDA: {tiling}")
+
+    elif torch.backends.mps.is_available():  # Check for Apple Silicon (MPS)
+        tile = {"x": 256, "y": 256, "z": 16}
+        halo = {"x": 16, "y": 16, "z": 4}
+        tiling = {"tile": tile, "halo": halo}
+        print(f"Determined tile size for MPS: {tiling}")
+
+    # I am not sure what is reasonable on a cpu. For now choosing very small tiling.
+    # (This will not work well on a CPU in any case.)
+    else:
+        tiling = {
+            "tile": {"x": 96, "y": 96, "z": 16},
+            "halo": {"x": 16, "y": 16, "z": 4},
+        }
+        print(f"Determining default tiling for CPU: {tiling}")
+
+    return tiling

flamingo_tools/plugin/segmentation_widget.py

@@ -9,8 +9,7 @@
 from qtpy.QtWidgets import QWidget, QVBoxLayout, QPushButton, QLabel, QComboBox
 
 from .base_widget import BaseWidget
-from .util import get_default_tiling, get_device
-from ..model_utils import get_model, get_model_registry
+from ..model_utils import get_model, get_model_registry, get_device, get_default_tiling
 
 
 def _load_custom_model(model_path: str, device: Optional[Union[str, torch.device]] = None) -> torch.nn.Module:

flamingo_tools/plugin/util.py

@@ -1,8 +0,0 @@
-
-
-def get_default_tiling():
-    pass
-
-
-def get_device():
-    pass

flamingo_tools/test_data.py

@@ -2,10 +2,12 @@
 from typing import Tuple
 
 import imageio.v3 as imageio
+import pooch
 import requests
 from skimage.data import binary_blobs, cells3d
 from skimage.measure import label
 
+from .file_utils import get_cache_dir
 from .segmentation.postprocessing import compute_table_on_the_fly
 
 SEGMENTATION_URL = "https://owncloud.gwdg.de/index.php/s/kwoGRYiJRRrswgw/download"
@@ -93,13 +95,43 @@ def create_test_data(root: str, size: int = 256, n_channels: int = 2, n_tiles: i
             imageio.imwrite(out_path, data)
 
 
+def _sample_registry():
+    urls = {
+        "PV": "https://owncloud.gwdg.de/index.php/s/JVZCOpkILT70sdv/download",
+        "VGlut3": "https://owncloud.gwdg.de/index.php/s/LvGXh0xQR9IKvNk/download",
+        "CTBP2": "https://owncloud.gwdg.de/index.php/s/qaffCaF1sGpqlT3/download",
+    }
+    registry = {
+        "PV": "fbf50cc9119f2dd2bd4dac7d76b746b7d42cab33b94b21f8df304478dd51e632",
+        "VGlut3": "6a3af6ffce3d06588ffdc73df356ac64b83b53aaf6aabeabd49ef6d11d927e20",
+        "CTBP2": "8dcd5f1ebb35194f328788594e275f2452de0e28c85073578dac7100d83c45fc",
+    }
+    cache_dir = get_cache_dir()
+    data_registry = pooch.create(
+        path=os.path.join(cache_dir, "data"),
+        base_url="",
+        registry=registry,
+        urls=urls,
+    )
+    return data_registry
+
+
 def sample_data_pv():
-    pass
+    data_path = _sample_registry().fetch("PV")
+    data = imageio.imread(data_path, extension=".tif")
+    add_image_kwargs = {"name": "PV", "colormap": "gray"}
+    return [(data, add_image_kwargs)]
 
 
 def sample_data_vglut3():
-    pass
+    data_path = _sample_registry().fetch("VGlut3")
+    data = imageio.imread(data_path, extension=".tif")
+    add_image_kwargs = {"name": "VGlut3", "colormap": "gray"}
+    return [(data, add_image_kwargs)]
 
 
 def sample_data_ctbp2():
-    pass
+    data_path = _sample_registry().fetch("CTBP2")
+    data = imageio.imread(data_path, extension=".tif")
+    add_image_kwargs = {"name": "CTBP2", "colormap": "gray"}
+    return [(data, add_image_kwargs)]