Pull from main

.github/workflows/documentation.yml

@@ -36,7 +36,8 @@ jobs:
 
     # Deploy the newly generated html pages to the GitHub Pages associated with this repo
     - name: Deploy to GitHub Pages
+      if: github.repository == 'brysongray/neurotrack'
       uses: peaceiris/actions-gh-pages@v4
       with:
         github_token: ${{ secrets.GITHUB_TOKEN }}
-        publish_dir: docs/build/html
+        publish_dir: docs/build/html

.gitignore

@@ -135,3 +135,16 @@ dmypy.json
 # Training data and model weights
 *.pt
 *.csv
+*.tif
+*.swc
+
+*.html
+*.pdf
+*.png
+*.jpg
+*.jpeg
+*.gif
+*.zip
+*.out
+*.json
+*.ipynb

README.md

@@ -33,4 +33,13 @@ No non-standard hardware is required, but this library uses pytorch which can us
 
 Examples for usage are given in Jupyter notebooks in the notebooks folder of the github repository.
 
+## Example neuron tracking results
+
+Gray surface is the true neuron, red surface is the result of automated tracing.
+
+![hippo](https://media.giphy.com/media/ZEcTAh6nwSrbRMdMDS/giphy.gif)
+
+![hippo](https://media.giphy.com/media/7Vu79HIrW4XROOVX03/giphy.gif)
+
+![hippo](https://media.giphy.com/media/Qlkpz2BHlb1bjyHEmo/giphy.gif)
 

bin/classifier_train.py

@@ -3,12 +3,15 @@
 """
 
 import argparse
+from glob import glob
 import os
 from pathlib import Path
 import sys
 import torch
-
-sys.path.append(str(Path(__file__).parents[1]))
+from torch.utils.data import DataLoader
+script_path = Path(os.path.abspath(__file__))
+parent_dir = script_path.parent.parent  # Go up two levels
+sys.path.append(str(parent_dir))
 import models
 from solvers import branch_classifier
 
@@ -25,19 +28,24 @@ def main():
     -o, --out: Path to output directory.
     -l, --learning_rate: Optimizer learning rate.
     -N, --epochs: Number of training epochs.
+    -w --weights: pretrained model weights 
     """
     parser = argparse.ArgumentParser()
     parser.add_argument('-s', '--source', type=str, help='Source directory containing labels as csv files and input images folder (observations).')
     parser.add_argument('-o','--out', type=str, help="Path to output directory.")
+    parser.add_argument('-c', '--channels', type=int, default=3, help='Number of input channels in the images.')
     parser.add_argument('-l', '--learning_rate', type=float, default=0.001, help='Optimizer learning rate.')
     parser.add_argument('-N', '--epochs', type=int, default=15, help='Number of training epochs.')
+    parser.add_argument('-w', '--weights', type=str, default=None, help='pretrained model weights')
 
     args = parser.parse_args()
     source = args.source
     out_dir = args.out
     lr = args.learning_rate
     epochs = args.epochs
-
+    weights = args.weights
+    channels = args.channels
+
     source_list = os.listdir(source)
     training_labels_file = [f for f in source_list if 'training_labels' in f][0]
     training_labels_file = os.path.join(source, training_labels_file)
@@ -49,30 +57,91 @@ def main():
     if not os.path.exists(test_labels_file):
         raise FileNotFoundError("Source directory must contain a csv file with `test_labels` in the filename,\
                                 but none was found")
-    img_dir = os.path.join(source, 'observations')
+    img_dir = [d for d in source_list if 'observations' in d][0]
+    img_dir = os.path.join(source, img_dir)
     if not os.path.exists(img_dir):
         raise FileNotFoundError("Source directory must contain a folder named `observations`,\
                                 but none was found.")
 
     transform = branch_classifier.transform # random permutation and flip
     training_data = branch_classifier.StateData(labels_file=training_labels_file,
-                            img_dir=img_dir,
-                            transform=transform)
+                            img_dir=img_dir,transform=transform)
     test_data = branch_classifier.StateData(labels_file=test_labels_file,
                             img_dir=img_dir)
 
-    training_dataloader = branch_classifier.init_dataloader(training_data)
-    test_dataloader = branch_classifier.init_dataloader(test_data)
+    training_dataloader = branch_classifier.init_dataloader(training_data, batchsize=64)
+    test_dataloader = branch_classifier.init_dataloader(test_data, batchsize=64)
 
     if not os.path.exists(out_dir):
         os.makedirs(out_dir, exist_ok=True)
 
-    classifier = models.ResNet(models.ResidualBlock, [3, 4, 6, 3], num_classes=1)
+    # classifier = models.ResNet2D(models.ResidualBlock2D, [3, 4, 6, 3], in_channels=36, num_classes=1)
+    classifier = models.ResNet3D(models.ResidualBlock3D, [3, 4, 6, 3], in_channels=channels, num_classes=1)
     classifier = classifier.to(device=DEVICE, dtype=dtype)
 
+    if weights is not None:
+        state_dict = torch.load(weights, map_location=DEVICE)
+        classifier.load_state_dict(state_dict)
+        print(f"Loaded pretrained weights from {weights}")
+
     branch_classifier.train(training_dataloader, test_dataloader, out_dir, lr, epochs, classifier, state_dict=None)
 
     return
 
+
+# # TODO: Change to online data sampling from swc and image files
+# def main():
+#     """
+#     Main function to train the branch classifier model.
+
+#     Arguments
+#     ---------
+#     -s, --swc_dir: Source directory containing the neuron trees in swc format.
+#     -i, --image_dir: Source directory containing the input images.
+#     -o, --out: Path to output directory.
+#     -l, --learning_rate: Optimizer learning rate.
+#     -N, --epochs: Number of training epochs.
+#     -w, --weights: pretrained model weights 
+#     """
+#     parser = argparse.ArgumentParser()
+#     parser.add_argument('-s', '--swc_dir', type=str, help='Source directory containing the neuron trees in swc format.')
+#     parser.add_argument('-i', '--image_dir', type=str, help='Source directory containing the input images.')
+#     parser.add_argument('-o', '--out', type=str, help="Path to output directory.")
+#     parser.add_argument('-l', '--learning_rate', type=float, default=0.001, help='Optimizer learning rate.')
+#     parser.add_argument('-N', '--epochs', type=int, default=15, help='Number of training epochs.')
+#     parser.add_argument('-w', '--weights', type=str, default=None, help='pretrained model weights')
+#     args = parser.parse_args()
+#     swc_dir = args.swc_dir
+#     image_dir = args.image_dir
+#     out_dir = args.out
+#     lr = args.learning_rate
+#     epochs = args.epochs
+#     weights = args.weights
+#     if not os.path.exists(swc_dir):
+#         raise FileNotFoundError(f"Source directory {swc_dir} does not exist.")
+#     if not os.path.exists(image_dir):
+#         raise FileNotFoundError(f"Source directory {image_dir} does not exist.")
+#     if not os.path.exists(out_dir):
+#         os.makedirs(out_dir, exist_ok=True)
+#     if not os.path.exists(weights):
+#         raise FileNotFoundError(f"Pretrained weights file {weights} does not exist.")
+
+#     swc_files = [f for x in os.walk(swc_dir) for f in glob(os.path.join(x[0], "*.swc"))]
+#     img_files = [f for x in os.walk(image_dir) for f in glob(os.path.join(x[0], "*image.tif"))]
+
+#     classifier = models.ResNet3D(models.ResidualBlock3D, [3, 4, 6, 3], in_channels=3, num_classes=1)
+#     classifier = classifier.to(device=DEVICE, dtype=dtype)
+
+#     if weights is not None:
+#         state_dict = torch.load(weights, map_location=DEVICE)
+#         classifier.load_state_dict(state_dict)
+#         print(f"Loaded pretrained weights from {weights}")
+
+#     transform = branch_classifier.transform  # random permutation and flip
+#     branch_classifier.train(swc_files, img_files, out_dir, lr, epochs, classifier, transform=transform)
+
+
+
+
 if __name__ == "__main__":
     main()

bin/collect_branch_data.py

@@ -5,12 +5,13 @@
 import argparse
 from datetime import datetime
 from glob import glob
+import numpy as np
 import os
 from pathlib import Path
 import sys
 
 sys.path.append(str(Path(__file__).parents[1]))
-from data_prep import load, collect
+from data_prep import load, collect, tree
 
 DATE = datetime.now().strftime("%m-%d-%y")
 
@@ -42,20 +43,32 @@ def main():
     parser.add_argument('-o','--out', type=str, help="Path to output directory.")
     parser.add_argument('-n', '--name', type=str, help='Output filename base.')
     parser.add_argument('-a', '--adjust', action='store_true', default=False, help='Set to true if neuron coordinates were rescaled to draw images.')
+    parser.add_argument('-r', '--remove_soma', action='store_true', default=False, help='Remove edges near the soma from the swc.')
     parser.add_argument('--n_samples', type=int, default=100, help='Number of samples to collect from each image file.')
     args = parser.parse_args()
     labels_dir = args.labels
     image_dir = args.images
     out_dir = args.out
     name = args.name
     adjust = args.adjust
+    remove_soma = args.remove_soma
     samples_per_file = args.n_samples
 
     # get sample points from swc files
     files = [f for x in os.walk(labels_dir) for f in glob(os.path.join(x[0], '*.swc'))]
     swc_lists = []
     for f in files:
-        swc_lists.append(load.swc(f))
+        swc_list = load.swc(f)
+        swc_list = np.array(swc_list)
+        if adjust:
+            min = np.min(swc_list[:, 2:5], axis=0)
+            max = np.max(swc_list[:, 2:5], axis=0)
+            vol = np.prod(max - min)
+            scale = np.round((5e7 / vol)**(1/3)) # scale depends on the volume
+            swc_list[...,:3] = (swc_list[...,:3] - min) * scale + np.array([10.0, 10.0, 10.0])
+        if remove_soma:
+            swc_list, _ = tree.remove_soma(swc_list, max_radius=7.0)
+        swc_lists.append(swc_list)
     fnames = [f.split('/')[-1].split('.')[0] for f in files]
 
     sample_points = collect.swc_random_points(samples_per_file, swc_lists, fnames, adjust=adjust)