load weights by default

wahabk · Feb 6, 2025 · b8409cf · b8409cf
1 parent f096140
commit b8409cf
Show file tree

Hide file tree

Showing 8 changed files with 24 additions and 17 deletions.
diff --git a/Colloidoscope_tutorial.ipynb b/Colloidoscope_tutorial.ipynb
@@ -18,7 +18,7 @@
       "source": [
         "# **Colloidoscope!** - Colab tutorial 🖥\n",
         "\n",
-        "Colloidoscope uses deep learning models which are usually demanding on CPUs or GPUs\n",
+        "Colloidoscope uses deep learning models which are usually demanding on CPUs\n",
         "\n",
         "**Colab** notebooks can provide you with a free GPU to use for you to run your predictions, to make use of this please go to the top menu and click \n",
         "\n",
@@ -97,7 +97,7 @@
         "image = dc.read_tif(path)\n",
         "diameter=9\n",
         "\n",
-        "df = dc.detect(image, diameter=diameter, run_on='cpu')\n"
+        "df = dc.detect(image, diameter=diameter)\n"
       ]
     }
   ],

diff --git a/README.md b/README.md
@@ -9,7 +9,7 @@ A project to track colloids using confocal and deep learning.
 
 Colloidoscope is essentially a U-net trained on simulated data. This is designed to to improve tracking on very small colloids (~200nm) with a very bad point spread function, that is anisotropic and has high z blur.
 
-Pretrained weights and hyperparameter training history available upon request.
+Pretrained weights are included by default. Hyperparameter training history available upon request.
 
 # Installation
 
@@ -43,7 +43,7 @@ Reader = dc.ExploreLifReader(args, kwargs)
 Then simply use ```dc.detect()``` to track your image, this will return a pandas DataFrame just like TrackPy.
 
 ```Python
-df = dc.detect(image, weights_path='path/to/model_weights.pt')
+df = dc.detect(image)
 ```
 
 Please check `examples/` for scripts and jupyter notebooks.
@@ -95,3 +95,7 @@ If you want to launch the container on a custom hard drive use:
 Thanks to [Yushi Yang @yangyushi](https://github.com/yangyushi) for contributing the position simulations and helping me with the entire project in general.
 
 This project wraps explore_lif which is written by Mathieu Leocmach and C. Ybert.
+
+## Citation
+
+Incoming! 🚧
diff --git a/colloidoscope/attention_unet_202206.pt b/colloidoscope/attention_unet_202206.pt
diff --git a/colloidoscope/predict.py b/colloidoscope/predict.py
@@ -176,16 +176,22 @@ def detect(input_array:np.ndarray, diameter:Union[int, list], model:torch.nn.Mod
                 out_channels=1,
                 channels=[32, 64, 128],
                 strides=[2,2],
-                # act=params['activation'],
-                # norm=params["norm"],
                 padding='valid',
             )
 
 
     model = torch.nn.DataParallel(model, device_ids=None) # parallelise model
 
     if weights_path is None:
-        weights_path = Path(__file__).parent / "attention_unet_202302.pt"
+        model = monai.networks.nets.AttentionUnet(
+            spatial_dims=3,
+            in_channels=1,
+            out_channels=1,
+            channels=[32, 64, 128],
+            strides=[2,2],
+            padding='valid',
+        )
+        weights_path = Path(__file__).parent / "attention_unet_202206.pt"
         model_weights = torch.load(str(weights_path), map_location=device) # read trained weights
         model.load_state_dict(model_weights) # add weights to model
     elif weights_path != "preloaded":

diff --git a/examples/Colloidoscope_Colab_tutorial.ipynb b/examples/Colloidoscope_Colab_tutorial.ipynb
@@ -82,12 +82,11 @@
         "import colloidoscope as cd\n",
         "\n",
         "dc = cd.DeepColloid()\n",
-        "model_weights_path = 'path/to/model_weights.pt'\n",
         "path = 'path/to/image.tif'\n",
         "image = dc.read_tif(path)\n",
         "diameter=9\n",
         "\n",
-        "df = dc.detect(image, weights_path=model_weights_path, diameter=diameter)\n"
+        "df = dc.detect(image, diameter=diameter)\n"
       ]
     }
   ],

diff --git a/examples/detect.ipynb b/examples/detect.ipynb
@@ -26,7 +26,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "First read your image as usual and define the path to your model weights"
+    "First read your image as a numpy array"
    ]
   },
   {
@@ -35,7 +35,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model_weights_path = 'path/to/model_weights.pt'\n",
     "path = 'path/to/image.tif'\n",
     "image = dc.read_tif(path) "
    ]
@@ -63,7 +62,7 @@
    "source": [
     "Then finally use `dc.detect()` as shown below, define your diameter since `TrackPy` is used on the output to find the positions\n",
     "\n",
-    "This will return a `pandas` DataFrame just like Trackpy"
+    "This expects a numpy array image and will return a `pandas` DataFrame just like Trackpy"
    ]
   },
   {
@@ -72,7 +71,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "df = dc.detect(image, weights_path=model_weights_path, diameter=9)"
+    "df = dc.detect(image, diameter=9)"
    ]
   },
   {

diff --git a/examples/detect.py b/examples/detect.py
@@ -3,8 +3,7 @@
 if __name__ == "__main__":
     dc = cd.DeepColloid()
 
-    model_weights_path = 'path/to/model_weights.pt'
     path = 'path/to/image.tif'
     image = dc.read_tif(path) 
 
-    df = dc.detect(image, weights_path=model_weights_path, diameter=9)
+    df = dc.detect(image, diameter=9)
diff --git a/setup.py b/setup.py
@@ -2,14 +2,14 @@
 
 setup(
         name="colloidoscope", 
-        version='0.1.1',
+        version='0.2.0',
         author="Abdelwahab Kawafi",
         author_email="<[email protected]>",
         description='My PhD project to track colloids using confocal and deep learning.',
         packages=find_packages(),
         include_package_data=True,
         package_dir={'colloidoscope': 'colloidoscope'},
-        package_data={'colloidoscope': ['attention_unet_202302.pt']},
+        package_data={'colloidoscope': ['attention_unet_202302.pt', 'attention_unet_202206.pt']},
         install_requires=[
                 'numpy',
                 'torch',