updating readme and aind model metadata

README.md

@@ -1,3 +1,90 @@
 ---
 license: cc-by-4.0
 ---
+---
+license: mit
+language:
+- en
+tags:
+- segmentation
+- microscopy
+- tissue
+---
+# aind-brain-tissue-segmentation
+
+<p align="center">
+    <img src="https://allenneuraldynamics.github.io/assets/img/AIND_logo.png" width="400"/>
+<p>
+
+<p align="center">
+    <a href=""><b>AIND</b></a> &nbsp&nbsp ｜ &nbsp&nbsp 🖥️ <a href="https://github.com/AllenNeuralDynamics/aind-brain-segmentation">GitHub</a> &nbsp&nbsp  | &nbsp&nbsp🤗 <a href="https://huggingface.co/AllenInstitute">Hugging Face</a>&nbsp&nbsp | &nbsp&nbsp🤖 <a href="https://www.allenneuraldynamics.org/">Allen Institute for Neural Dynamics</a>&nbsp&nbsp
+<br>
+
+-----
+
+[**Brain Tissue Segmentation - Lightsheet Microscopy**]() <be>
+
+In this repository, we present a tool for fast brain tissue segmentation in lightsheet microscopy. This model was trained in entire 2D slices of data coming from our SmartSPIM Lightsheet Microscopes. We used a downsampled resolution corresponding to 14.4 microns in XY and 16.0 microns in Z that then were resized to 1024x1024.
+
+We also provide a script that perform fast inference in entire brains stored in OMEZarr format.
+
+
+## Brain Segmentation Examples
+
+<a href="https://github.com/user-attachments/assets/1021a3a3-cfa1-460c-b1d7-29746ecf764c">Video Sample 771602</a>
+
+<a href="https://github.com/user-attachments/assets/66ec0a36-6798-4081-98e4-b6e744355e99">Video Sample 782499</a>
+
+## 📑 Todo List
+- [ ] Multi-channel segmentation.
+- [ ] Generalization to other microscopes. E.g., (ExASPIM, TissueCyte)
+
+
+## Inference
+
+It supports two processing modes:
+1. **Lazy segmentation**: For large datasets that cannot fit into memory, images are resized to **1024x1024**, normalized, segmented, upsampled, and stored in a Zarr format.
+2. **In-memory segmentation**: For smaller datasets, a multiscale version is loaded into memory for processing.
+
+The model was trained with a resolution of **14.4 microns in XY** and **16 microns in Z**.
+
+<a href="https://github.com/AllenNeuralDynamics/aind-brain-segmentation/blob/main/code/large_scale_mask_prediction.py">Whole-brain segmentation script</a>
+
+
+```
+python large_scale_mask_prediction.py --input_path path/to/your/image.zarr
+```
+
+## Input Parameters
+| Parameter | Description |
+|-----------|-------------|
+| `image_path` | Path to the OME-Zarr dataset (local or S3). Example: `s3://bucket/SmartSPIM_Dataset/Ex_639_Em_667.zarr` |
+| `model_path` | Path to the trained segmentation model. |
+| `output_folder` | Directory where results will be saved. |
+| `target_size_mb` | Memory allocation target for large-scale predictions. |
+| `n_workers` | Number of workers for PyTorch DataLoader. |
+| `super_chunksize` | Data shard size for optimized cloud communication. |
+| `scale` | Name of the multiscale dataset to use for segmentation. |
+| `scratch_folder` | Path for temporary files (Raw data). |
+| `image_height` | Resize height for segmentation (Default: `1024`). |
+| `image_width` | Resize width for segmentation (Default: `1024`). |
+| `prob_threshold` | Probability threshold for segmentation mask generation. |
+
+## Outputs
+| Output File | Description |
+|-------------|-------------|
+| `probabilities.zarr` | Zarr dataset containing segmentation probabilities (useful for post-processing). |
+| `segmentation_mask.zarr` | Segmentation mask in the original image space (before resizing). |
+| `data.zarr` | Raw input data used for segmentation (optional). |
+
+#### Installation
+Clone the repo:
+```
+git clone https://github.com/AllenNeuralDynamics/aind-brain-segmentation
+cd aind-brain-segmentation
+```
+
+Install dependencies:
+```
+docker build -t brain_seg -f Dockerfile .
+```

model.json

@@ -0,0 +1,89 @@
+{
+   "describedBy": "https://raw.githubusercontent.com/AllenNeuralDynamics/aind-data-schema/main/src/aind_data_schema/core/model.py",
+   "schema_version": "0.0.1",
+   "name": "smartspim_tissue_segmentation.ckpt",
+   "license": "CC-BY-4.0",
+   "developer_full_name": [
+      "Camilo Laiton"
+   ],
+   "developer_institution": {
+      "name": "Allen Institute for Neural Dynamics",
+      "abbreviation": "AIND",
+      "registry": {
+         "name": "Research Organization Registry",
+         "abbreviation": "ROR"
+      },
+      "registry_identifier": "04szwah67"
+   },
+   "modality": [
+      {
+         "name": "Selective plane illumination microscopy",
+         "abbreviation": "SPIM"
+      }
+   ],
+   "architecture": {
+      "backbone": "UNet",
+      "software": [
+         {
+            "name": "lightning",
+            "version": "2.5.0.post0",
+            "url": null,
+            "parameters": {}
+         }
+      ],
+      "layers": null,
+      "parameters": {
+         "image_width": 1024,
+         "image_height": 1024
+      },
+      "notes": null
+   },
+   "intended_use": "2D tissue segmentation from SmartSPIM datasets",
+   "limitations": "Only trained on 639 channel using a resolution",
+   "pretrained_source_url": null,
+   "training": [
+      {
+         "name": "Model training",
+         "software_version": null,
+         "start_date_time": "2025-03-01T12:00:00-08:00",
+         "end_date_time": "2025-03-01T14:00:00-08:00",
+         "input_location": [
+            "s3://aind-msma-morphology-data/test_data/SmartSPIM/brain_segmentation_2d_allaxis_channel_639"
+         ],
+         "output_location": "Path to deployment",
+         "code_url": "https://github.com/AllenNeuralDynamics/aind-brain-segmentation",
+         "code_version": null,
+         "parameters": {
+            "learning_rate": 0.00001,
+            "batch_size": 32,
+            "augmentation": true
+         },
+         "outputs": {},
+         "notes": "Slices were pulled from a multiscaled image (14.4 microns in XY and 16.0 microns in Z) and then resized to 1024x1024",
+         "resources": null,
+         "train_performance": [
+            {
+               "name": "dice",
+               "value": 0.96
+            },
+            {
+               "name": "jaccard",
+               "value": 0.93
+            }
+         ],
+         "test_performance": [
+            {
+               "name": "precision",
+               "value": 0.95
+            },
+            {
+               "name": "recall",
+               "value": 0.91
+            }
+         ],
+         "test_data": "4:1 train/test split"
+      }
+   ],
+   "evaluations": [],
+   "notes": null
+}