feat: Huggingface

.gitignore

@@ -166,3 +166,12 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# Local scratch notes / codebase-state snapshots (not for version control)
+notes/
+
+# Downloaded Hugging Face checkpoint cache
+octopi/cache/
+
+# Claude markdown files
+CLAUDE.md

README.md

@@ -1,6 +1,6 @@
 # OCTOPI 🐙🐙🐙
 
-[![License](https://img.shields.io/pypi/l/octopi.svg?color=green)](https://github.com/chanzuckerberg/octopi/raw/main/LICENSE)
+[![License](https://img.shields.io/pypi/l/octopi.svg?color=green)](https://github.com/Biohub/octopi/raw/main/LICENSE)
 [![PyPI](https://img.shields.io/pypi/v/octopi.svg?color=green)](https://pypi.org/project/octopi)
 [![Python Version](https://img.shields.io/pypi/pyversions/octopi.svg?color=green)](https://www.python.org/)
 
@@ -47,7 +47,7 @@ octopi
 
 ## 📚 Documentation
 
-For detailed documentation, tutorials, CLI and API reference, visit our [documentation](https://chanzuckerberg.github.io/octopi/).
+For detailed documentation, tutorials, CLI and API reference, visit our [documentation](https://biohub.github.io/octopi/).
 
 ## 🤝 Contributing
 

docs/api/inference.md

@@ -7,9 +7,10 @@ This page covers running inference with trained octopi models, including segment
 Run trained models on tomograms to generate segmentation masks. The segmentation process takes your trained model weights and configuration to produce probability maps for each object class defined in your training targets.
 
 !!! info "Segmentation Parameters"
-    - **model_weights**: Path to your trained model weights (`.pth` file from training)
-    - **model_config**: Path to model configuration (`.yaml` file from training)
-    - **seg_info**: Tuple defining where to save segmentation results (`name`, `user_id`, `session_id`)
+    - **model_weights**: Path to your trained model weights (`.pth` file from training), or a pretrained checkpoint alias (e.g. `'tomogram-boundary'`) to auto-download from the [Hugging Face Hub](https://huggingface.co/biohub/octopi)
+    - **model_config**: Path to model configuration (`.yaml` file from training). Omit when `model_weights` is a checkpoint alias — its config is bundled and downloaded automatically
+    - **tomo_uri**: Tomogram URI in the form `"algorithm@voxel_size"` (e.g. `"denoised@10.0"`)
+    - **seg_uri**: Segmentation output URI in the form `"name:user_id/session_id"` (e.g. `"predict:octopi/1"`)
     - **use_tta**: Whether to use test-time augmentation for improved robustness
     - **run_ids**: Optional list of specific tomograms to process (None for all available)
 
@@ -27,7 +28,7 @@ The algorithm automatically applies size filtering based on object radii defined
 
     ```python
     import numpy as np
-    from octopi.pytorch.segmentation import Predictor
+    from octopi.pytorch.inference import Predictor
 
     model_weights = 'model_output/best_model.pth'
     model_config = 'model_output/model_config.yaml'
@@ -59,32 +60,42 @@ The algorithm automatically applies size filtering based on object radii defined
     config = 'eval_config.json'
     model_weights = 'model_output/best_model.pth'
     model_config = 'model_output/model_config.yaml'
-    seg_info = ['predict', 'octopi', '1']  # (name, user_id, session_id)
 
     segment(
         config=config,
-        tomo_algorithm='denoised',
-        voxel_size=10.0,
         model_weights=model_weights,
         model_config=model_config,
-        seg_info=seg_info,
+        tomo_uri='denoised@10.0',
+        seg_uri='predict:octopi/1',
         ntta=4
     )
     ```
 
+    You can also skip training entirely and segment with a pretrained checkpoint from the
+    [Hugging Face Hub](https://huggingface.co/biohub/octopi) — the weights and matching
+    config are downloaded and cached automatically, so `model_config` can be omitted:
+
+    ```python
+    segment(
+        config=config,
+        model_weights='tomogram-boundary',
+        tomo_uri='denoised@10.0',
+        seg_uri='predict:octopi/1',
+    )
+    ```
+
     <details markdown="1">
     <summary><strong>💡 segment() reference</strong></summary>
 
-    `segment(config, tomo_algorithm, voxel_size, model_weights, model_config, seg_info=['predict', 'octopi', '1'], run_ids=None, batch_size=1, ntta=4)`
+    `segment(config, model_weights, model_config=None, tomo_uri='wbp@10.0', seg_uri='predict:octopi/1', run_ids=None, batch_size=1, swbs=4, overlap=0.5, ntta=4)`
 
     **Parameters:**
 
     - `config` (str): Path to CoPick configuration file
-    - `tomo_algorithm` (str): Tomogram algorithm identifier
-    - `voxel_size` (float): Voxel spacing in Angstroms
-    - `model_weights` (str or list): Path(s) to trained model weights (.pth file(s))
-    - `model_config` (str or list): Path(s) to model configuration (.yaml file(s))
-    - `seg_info` (list): Output segmentation identifier `(name, user_id, session_id)`
+    - `model_weights` (str or list): Path(s) to trained model weights (.pth file(s)), or a pretrained checkpoint alias (e.g. `'tomogram-boundary'`)
+    - `model_config` (str or list): Path(s) to model configuration (.yaml file(s)). May be omitted when `model_weights` is a checkpoint alias
+    - `tomo_uri` (str): Tomogram URI in the form `"algorithm@voxel_size"` (default: `'wbp@10.0'`)
+    - `seg_uri` (str): Segmentation output URI in the form `"name:user_id/session_id"` (default: `'predict:octopi/1'`)
     - `run_ids` (list): Specific run IDs to process (default: None — all runs)
     - `batch_size` (int): Tomograms processed concurrently per GPU (default: 1)
     - `ntta` (int): Number of test-time augmentation rotations (default: 4; set to 1 to disable)

docs/api/quick-start.md

@@ -91,11 +91,10 @@ from octopi.workflows import segment
 
 segment(
     config=config,
-    tomo_algorithm='denoised',
-    voxel_size=10.012,
     model_weights=f'{results_folder}/best_model.pth',
     model_config=f'{results_folder}/model_config.yaml',
-    seg_info=['predict', 'octopi', '1'],
+    tomo_uri='denoised@10.012',
+    seg_uri='predict:octopi/1',
     ntta=4  # number of test-time augmentation rotations
 )
 ```
@@ -200,11 +199,10 @@ evaluate(
     print("Step 3: Running segmentation...")
     segment(
         config=config,
-        tomo_algorithm=tomo_algorithm,
-        voxel_size=voxel_size,
         model_weights=f'{results_folder}/best_model.pth',
         model_config=f'{results_folder}/model_config.yaml',
-        seg_info=['predict', 'octopi', '1'],
+        tomo_uri=f'{tomo_algorithm}@{voxel_size}',
+        seg_uri='predict:octopi/1',
         ntta=4
     )
 

docs/getting-started/installation.md

@@ -14,7 +14,7 @@ This will install the latest stable release along with all required dependencies
 If you want to contribute to octopi or need the latest development version, you can install from source:
 
 ```bash
-git clone https://github.com/chanzuckerberg/octopi.git
+git clone https://github.com/Biohub/octopi.git
 cd octopi
 pip install -e .
 ```
@@ -50,12 +50,12 @@ The editable (-e) install ensures that local code changes are immediately reflec
     │ model-explore   Perform model architecture search with Optuna.                  │
     ╰─────────────────────────────────────────────────────────────────────────────────╯
     ╭─ Inference ─────────────────────────────────────────────────────────────────────╮
-    │ segment           Segment volumes using trained neural network models.          │
-    │ localize          Convert Segmentation Masks to 3D Particle Coordinates.        │
-    │ membrane-extract  Extract membrane-bound picks based on proximity to organelle  │
-    │                   or membrane segmentation.                                     │
-    │ evaluate          Evaluate particle localization performance against ground     │
-    │                   truth annotations.                                            │
+    │ segment    Segment volumes using trained neural network models.                 │
+    │ localize   Convert Segmentation Masks to 3D Particle Coordinates.               │
+    │ extract    Extract objects or picks from existing pipeline outputs (isolate a   │
+    │            single object's mask, or split picks by membrane proximity).         │
+    │ evaluate   Evaluate particle localization performance against ground truth      │
+    │            annotations.                                                         │
     ╰─────────────────────────────────────────────────────────────────────────────────╯
     ```
 

docs/getting-started/quickstart.md

@@ -17,10 +17,9 @@ Create semantic masks for your proteins of interest using annotation metadata:
 ```bash
 octopi create-targets \
     --config config.json \
-    --tomo-alg wbp --voxel-size 10 \
+    --tomo-uri wbp@10.0 \
     --picks-user-id data-portal --picks-session-id 0 \
-    --target-session-id 1 --target-segmentation-name targets \
-    --target-user-id octopi
+    --target-uri targets:octopi/1
 ```
 
 🎯 This creates training targets for a single copick query. To produce targets from multiple coordinate queries,  refer to the [Prepare Labels](../user-guide/labels.md) section.

docs/index.md

@@ -96,4 +96,4 @@
 
 ## Getting Help
 
-Open an issue on our [GitHub repository](https://github.com/chanzuckerberg/octopi).
+Open an issue on our [GitHub repository](https://github.com/Biohub/octopi).

docs/user-guide/claude-code-mcp.md

@@ -32,7 +32,7 @@ Once connected, you can describe what you want in plain language and Claude hand
 
     ---
 
-    Executes `create-targets`, `localize`, `evaluate`, and `membrane-extract` directly and reports back.
+    Executes `create-targets`, `localize`, `evaluate`, and `extract` directly and reports back.
 
 -   :material-chip:{ .lg .middle } **Hand off GPU jobs**
 
@@ -126,7 +126,7 @@ octopi create-targets \
     --target ribosome,manual,1 \
     --target virus-like-particle,tm,2 \
     --target membranes,membrane-seg,1 \
-    --voxel-size 10
+    --tomo-uri wbp@10.0
 ```
 
 ??? tip "URI shorthand for tomograms, segmentations, and picks"

docs/user-guide/data-import.md

@@ -78,9 +78,7 @@ Copick provides a CLI for generating configuration files that mount either local
         --proj-description "Synaptic Vesicles collected on 24sep24"
     ```
 
-    We can define either objects that are continuous segmentations (e.g., organelles or memebranes) or coordinates. For pickable objects, we can store meta-data including the particle radius and a corresponding PDB-ID:
-
-        - `--objects name,is_particle,radius,pdb_id`.
+    We can define either objects that are continuous segmentations (e.g., organelles or memebranes) or coordinates. For pickable objects, we can store meta-data including the particle radius and a corresponding PDB-ID: `--objects name,is_particle,radius,pdb_id`.
 
     ??? info "`copick config filesystem` parameters"
 
@@ -124,6 +122,7 @@ If you have tomograms stored locally in `*.mrc` format (e.g., from Warp, IMOD, o
 
     ```bash 
     copick add tomogram \
+        ".mrc" \
         -c /path/to/config.json \
         --tomo-type sart \
     ```

docs/user-guide/inference.md

@@ -8,7 +8,8 @@ Octopi inference follows a systematic two-step approach:
 
 1. **Segmentation** - Apply trained model to generate 3D probability masks with test-time augmentation (TTA).
 2. **Localization** - Convert probability masks into 3D coordinates using size-based filtering.
-3. **Evaluation (Optional)** - Compare predicted coordinates against ground truth annotations.
+3. **Extraction (Optional)** - Isolate a single object from a prediction, or split picks by membrane proximity.
+4. **Evaluation (Optional)** - Compare predicted coordinates against ground truth annotations.
 
 ??? tip "Parallelism and Resource Utilization"
 
@@ -53,8 +54,8 @@ octopi segment \
 
         | Parameter | Description | Notes |
         |----------|-------------|------|
-        | `--model-config` | Model configuration file(s). | Required; comma-separated for ensembles |
-        | `--model-weights` | Model weight file(s). | Must match `--model-config` order |
+        | `--model-config` | Model configuration file(s). | Required for local weights; comma-separated for ensembles. Omit when `--model-weights` is a Hugging Face checkpoint alias |
+        | `--model-weights` | Model weight file(s), or a pretrained checkpoint alias. | Must match `--model-config` order for local files |
 
     === "Inference"
 
@@ -76,6 +77,24 @@ octopi segment \
     --seg-uri ensemble:octopi/1
 ```
 
+### Pretrained Checkpoints
+
+In development! Only one pre-trained model is available. More to come.
+
+!!! tip "Skip training with a pretrained checkpoint"
+
+    Pass a checkpoint name from the  Hugging Face Hub repo as `--model-weights`. It's downloaded and cached automatically, so `--model-config` can be omitted.
+
+    ```bash
+    octopi segment \
+        --config config.json \
+        --tomo-uri wbp@10.0 \
+        --model-weights tomogram-boundary \
+        --seg-uri predict:octopi/1
+    ```
+
+    💡 - See the [model card](https://huggingface.co/biohub/octopi) for the full list of available checkpoints.
+
 ---
 
 ## Localization 
@@ -120,13 +139,58 @@ The localization algorithm uses **particle size information** from your copick c
         | `--pick-session-id` | Session ID for particle picks. | `1` | Used for result grouping |
         | `--pick-user-id` | User ID for particle picks. | `octopi` | Used for result grouping |
 
-## Context-Aware Particle Extraction
+## Extraction (Optional)
+
+`octopi extract` is a command group for post-processing outputs you've already generated — isolating a single object from a raw multi-class prediction, or splitting picks by proximity to a membrane/organelle segmentation.
 
-This optional post-processing step splits an existing set of particle picks into two groups:
+### Isolate a Single Object
+
+Pull one object's mask out of a raw multi-class `octopi segment` prediction and save it as its own standalone segmentation.
+
+```bash
+octopi extract seg \
+    --config config.json \
+    --seg-uri predict:octopi/1 \
+    --name membranes \
+    --session-id 1
+```
+
+Octopi reads the inference log written by `octopi segment` to recover the voxel size and the object's integer label within the raw prediction, then writes a binary mask for just that object as a new segmentation.
+
+??? info "`octopi extract seg -h`"
+
+    === "Input"
+
+        | Parameter | Description | Default | Notes |
+        |----------|-------------|---------|------|
+        | `--config` | Path to the CoPick configuration file. | – | Required |
+        | `--name` | Object name to extract from the raw multi-class prediction. | – | Required. Example: `membranes` |
+        | `--seg-uri` | Source segmentation to extract from (`name:user_id/session_id`). | `predict:octopi/1` | Must be a prediction written by `octopi segment` |
+        | `--run-ids` | Specific run IDs to process. | All runs | Example: `run1,run2` |
+
+    === "Output"
+
+        | Parameter | Description | Default | Notes |
+        |----------|-------------|---------|------|
+        | `--user-id` | User ID for the extracted segmentation. | Source segmentation's user ID | |
+        | `--session-id` | Session ID for the extracted segmentation. | `1` | |
+
+### Membrane-Proximity Picks
+
+This step splits an existing set of particle picks into two groups:
 
 - **Membrane-close picks**: particles within a configurable distance threshold of a membrane/organelle segmentation.
 - **Membrane-far picks**: particles outside that threshold.
 
+```bash
+octopi extract mb-picks \
+    --config config.json \
+    --picks-uri ribosome:octopi/1 \
+    --seg-uri membrane:membrain-seg/1 \
+    --save-user-id octopi \
+    --save-session-id 10
+```
+
 For membrane-close particles, we can also **align orientations** so that each particle’s rotation is consistent with the local membrane normal (estimated from the vector between the particle and the closest organelle center).
 
 ??? tip "What this is useful for"
@@ -240,17 +304,6 @@ octopi evaluate
 
 ---
 
-## Visualization
-
-To visualize your results and validate the quality of segmentations and coordinates, refer to our interactive notebook:
-
-**📓 [Inference Notebook](https://github.com/chanzuckerberg/octopi/blob/main/notebooks/inference.ipynb)**
-
-With this notebook, we can overlay the segmentation masks or coordiantes the tomograms. 
-
-![Coordinates Visualization](../assets/coordinates.png)
-*Example of predicted particle coordinates displayed on a holdout tomogram from cryo-ET training dataset. The visualization shows Octopi's localization results overlaid on tomographic data from [DatasetID: 10440](https://cryoetdataportal.czscience.com/datasets/10440).*
-
 ## Next Steps
 
 You now have a complete workflow for applying Octopi models to new tomographic data. The inference pipeline transforms your trained models into actionable scientific results through robust segmentation, intelligent localization, and comprehensive evaluation.