octopi 1.0__tar.gz

octopi-1.0/LICENSE

@@ -0,0 +1,41 @@
+# Legal
+
+## License for the octopi package
+
+This package is licensed under the MIT License:
+
+```
+MIT License
+
+Copyright (c) 2025 Chan Zuckerberg Initiative
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+```
+
+## License Notice for Dependencies
+
+```
+This repository is licensed under the MIT License; however, it relies on certain third-party dependencies that are licensed under the GNU General Public License (GPL). Specifically:
+
+- monai is licensed under the Apache License 2.0.
+- pytorch-lightning is licensed under the Apache License 2.0.
+
+All dependencies use permissive open-source licenses that are compatible with this project's MIT License. No GPL or other copyleft licensed dependencies are included.
+For specific licensing information about any dependency, please refer to the respective package documentation or repository.
+```

octopi-1.0/PKG-INFO

@@ -0,0 +1,209 @@
+Metadata-Version: 2.3
+Name: octopi
+Version: 1.0
+Summary: Model architecture exploration for cryoET particle picking
+License: MIT
+Author: Jonathan Schwartz
+Requires-Python: >=3.9,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: copick
+Requires-Dist: ipywidgets
+Requires-Dist: kaleido
+Requires-Dist: matplotlib
+Requires-Dist: mlflow (==2.17.0)
+Requires-Dist: monai-weekly (==1.5.dev2448)
+Requires-Dist: mrcfile
+Requires-Dist: multiprocess
+Requires-Dist: nibabel
+Requires-Dist: optuna (==4.0.0)
+Requires-Dist: optuna-integration[botorch,pytorch-lightning]
+Requires-Dist: pandas
+Requires-Dist: plotly
+Requires-Dist: python-dotenv
+Requires-Dist: pytorch-lightning (==2.4.0)
+Requires-Dist: requests (>=2.25.1,<3.0.0)
+Requires-Dist: seaborn
+Requires-Dist: torch-ema
+Requires-Dist: tqdm
+Description-Content-Type: text/markdown
+
+# OCTOPI 🐙🐙🐙
+**O**bject dete**CT**ion **O**f **P**rote**I**ns. A deep learning framework for Cryo-ET 3D particle picking with autonomous model exploration capabilities.
+
+## 🚀 Introduction
+
+octopi addresses a critical bottleneck in cryo-electron tomography (cryo-ET) research: the efficient identification and extraction of proteins within complex cellular environments. As advances in cryo-ET enable the collection of thousands of tomograms, the need for automated, accurate particle picking has become increasingly urgent.
+
+Our deep learning-based pipeline streamlines the training and execution of 3D autoencoder models specifically designed for cryo-ET particle picking. Built on [copick](https://github.com/copick/copick), a storage-agnostic API, octopi seamlessly accesses tomograms and segmentations across local and remote environments. 
+
+## 🧩 Features
+
+octopi offers a modular, deep learning-driven pipeline for:
+*	Training and evaluating custom 3D U-Net models for particle segmentation.
+*	Automatically exploring model architectures using Bayesian optimization via Optuna.
+*	Performing inference for both semantic segmentation and particle localization.
+
+octopi empowers researchers to navigate the dense, intricate landscapes of cryo-ET datasets with unprecedented precision and efficiency without manual trial and error.
+
+## Getting Started
+### Installation
+
+*Octopi* is available on PyPI.
+```
+pip install octopi
+```
+
+## 📚 Usage
+
+octopi provides a clean, scriptable command-line interface. Run the following command to view all available subcommands:
+```
+octopi --help
+```
+Each subcommand supports its own --help flag for detailed usage. To see practical examples of how to interface directly with the octopi API, explore the notebooks/ folder.
+
+If you're running octopi on an HPC cluster, several SLURM-compatible submission commands are available. You can view them by running:
+```
+octopi-slurm --help
+```
+This provides utilities for submitting training, inference, and localization jobs in SLURM-based environments.
+
+### 📥 Data Import & Preprocessing
+
+To train or run inference with octopi, your tomograms must be organized inside a CoPick project. octopi supports two primary methods for data ingestion, both of which include optional Fourier cropping to reduce resolution and accelerate downstream processing.
+
+If your tomograms are already processed and stored locally in .mrc format (e.g., from Warp, IMOD, or AreTomo), you can import them into a new or existing CoPick project using:
+
+```
+octopi import-mrc-volumes \
+    --input-folder /path/to/mrc/files --config /path/to/config.json \
+    --target-tomo-type denoised --input-voxel-size --output-voxel-size 10
+```
+
+octopi also can process tomograms that are hosted on the data portal. Users can download tomograms onto their own remote machine especially if they would like to downsample the tomograms to a lower resolution for speed and memory. You can download and process the tomograms using:
+```
+octopi download-dataportal \
+    --config /path/to/config.json --datasetID 10445 --overlay-path path/to/saved/zarrs \
+    --input-voxel-size 5 --output-voxel-size 10 \
+    --dataportal-name wbp --target-tomotype wbp 
+```
+
+### 📁 Training Labels Preparation  
+
+Use `octopi create-targets` to create semantic masks for proteins of interest using annotation metadata. In this example lets generate picks segmentations for dataset 10439 from the CZ cryoET Dataportal (only need to run this step once). 
+```
+octopi create-targets \
+    --config config.json \
+    --target apoferritin --target beta-galactosidase,slabpick,1 \
+    --target ribosome,pytom,0 --target virus-like-particle,pytom,0 \
+    --seg-target membrane \
+    --tomo-alg wbp --voxel-size 10 \
+    --target-session-id 1 --target-segmentation-name remotetargets \
+    --target-user-id train-octopi
+```
+
+### 🧠 Training a single 3D U-Net model  
+Train a 3D U-Net model on the prepared datasets using the prepared target segmentations. We can use tomograms derived from multiple copick projects. 
+``` 
+octopi train-model \
+    --config experiment,config1.json \
+    --config simulation,config2.json \
+    --voxel-size 10 --tomo-alg wbp --Nclass 8 \
+    --tomo-batch-size 50 --num-epochs 100 --val-interval 10 \
+    --target-info remotetargets,train-octopi,1
+```
+Outputs will include model weights (.pth), logs, and training metrics.
+
+### 🔍 Model exploration with Optuna
+
+octopi🐙 supports automatic neural architecture search using Optuna, enabling efficient discovery of optimal 3D U-Net configurations through Bayesian optimization. This allows users to maximize segmentation accuracy without manual tuning.
+
+To launch a model exploration job:
+```
+octopi model-explore \
+    --config experiment,/mnt/dataportal/ml_challenge/config.json \
+    --config simulation,/mnt/dataportal/synthetic_ml_challenge/config.json \
+    --voxel-size 10 --tomo-alg wbp --Nclass 8 \
+    --model-save-path train_results
+```
+Each trial evaluates a different architecture and logs:
+	•	Segmentation performance metrics
+	•	Model weights and configs
+	•	Training curves and validation loss
+
+🔬 Trials are automatically tracked with MLflow and saved under the specified `--model-save-path`.
+
+#### Optuna Dashboard
+
+To quickly asses the exploration results and observe which trials results the best architectures, Optuna provides a dashboard that summarizes all the information on a dashboard. The instrucutions to access the dashboard are available here - https://optuna-dashboard.readthedocs.io/en/latest/getting-started.html, it is recommended to use either VS-Code extension or CLI. 
+ 
+#### 📊 MLflow experiment tracking   
+
+To use CZI cloud MLflow tracker, add a `.env` in the root directory like below. You can get a CZI MLflow access token from [here](https://mlflow.cw.use4-prod.si.czi.technology/api/2.0/mlflow/users/access-token) (note that a new token will be generated everytime you open this site).
+```
+MLFLOW_TRACKING_USERNAME = <Your_CZ_email>
+MLFLOW_TRACKING_PASSWORD = <Your_mlflow_access_token>
+```
+
+octopi supports MLflow for logging and visualizing model training and hyperparameter search results, including:
+	•	Training loss/validation metrics over time
+	•	Model hyperparameters and architecture details
+	•	Trial comparison (e.g., best performing model)
+
+You can use either a local MLflow instance, a remote (HPC) instance, or the CZI cloud server:
+
+#### 🧪 Local MLflow Dashboard
+
+To inspect results locally: `mlflow ui` and open http://localhost:5000 in your browser.
+
+#### 🖥️ HPC Cluster MLflow Access (Remote via SSH tunnel)
+
+If running octopi on a remote cluster (e.g., Biohub Bruno), forward the MLflow port. 
+On your local machine: 
+ `ssh -L 5000:localhost:5000 remote_username@remote_host` (in the case of Bruno the remote would be `login01.czbiohub.org`). 
+ 
+Then on the remote terminal (login node): ` mlflow ui --host 0.0.0.0 --port 5000` to launch the MLFlow dashboard on a local borwser.
+
+#### ☁️ CZI coreweave cluser
+
+For the CZI coreweave cluser, MLflow is already hosted. Go to the CZI [mlflow server](https://mlflow.cw.use4-prod.si.czi.technology/). 
+
+🔐 A .env file is required to authenticate (see Getting Started section).
+📁 Be sure to register your project name in MLflow before launching runs.
+
+### 🔮 Segmentation
+Generate segmentation prediction masks for tomograms in a given copick project.
+```
+octopi inference \
+    --config config.json \
+    --seg-info predict,unet,1 \
+    --model-config train_results/best_model_config.yaml \
+    --model-weights train_results/best_model.pth \
+    --voxel-size 10 --tomo-alg wbp --tomo-batch-size 25
+```
+Output masks will be saved to the corresponding copick project under the `seg-info` input.
+
+### 📍 Localization
+Convert the segmentation masks into particle coordinates. 
+```
+octopi localize \
+    --config config.json \
+    --pick-session-id 1 --pick-user-id unet \
+    --seg-info predict,unet,1
+```
+
+## Contributing
+
+This project adheres to the Contributor Covenant code of conduct. By participating, you are expected to uphold this code. Please report unacceptable behavior to opensource@chanzuckerberg.com.
+
+## Reporting Security Issues
+
+Please note: If you believe you have found a security issue, please responsibly disclose by contacting us at security@chanzuckerberg.com.
+
+
+

octopi-1.0/README.md

@@ -0,0 +1,173 @@
+# OCTOPI 🐙🐙🐙
+**O**bject dete**CT**ion **O**f **P**rote**I**ns. A deep learning framework for Cryo-ET 3D particle picking with autonomous model exploration capabilities.
+
+## 🚀 Introduction
+
+octopi addresses a critical bottleneck in cryo-electron tomography (cryo-ET) research: the efficient identification and extraction of proteins within complex cellular environments. As advances in cryo-ET enable the collection of thousands of tomograms, the need for automated, accurate particle picking has become increasingly urgent.
+
+Our deep learning-based pipeline streamlines the training and execution of 3D autoencoder models specifically designed for cryo-ET particle picking. Built on [copick](https://github.com/copick/copick), a storage-agnostic API, octopi seamlessly accesses tomograms and segmentations across local and remote environments. 
+
+## 🧩 Features
+
+octopi offers a modular, deep learning-driven pipeline for:
+*	Training and evaluating custom 3D U-Net models for particle segmentation.
+*	Automatically exploring model architectures using Bayesian optimization via Optuna.
+*	Performing inference for both semantic segmentation and particle localization.
+
+octopi empowers researchers to navigate the dense, intricate landscapes of cryo-ET datasets with unprecedented precision and efficiency without manual trial and error.
+
+## Getting Started
+### Installation
+
+*Octopi* is available on PyPI.
+```
+pip install octopi
+```
+
+## 📚 Usage
+
+octopi provides a clean, scriptable command-line interface. Run the following command to view all available subcommands:
+```
+octopi --help
+```
+Each subcommand supports its own --help flag for detailed usage. To see practical examples of how to interface directly with the octopi API, explore the notebooks/ folder.
+
+If you're running octopi on an HPC cluster, several SLURM-compatible submission commands are available. You can view them by running:
+```
+octopi-slurm --help
+```
+This provides utilities for submitting training, inference, and localization jobs in SLURM-based environments.
+
+### 📥 Data Import & Preprocessing
+
+To train or run inference with octopi, your tomograms must be organized inside a CoPick project. octopi supports two primary methods for data ingestion, both of which include optional Fourier cropping to reduce resolution and accelerate downstream processing.
+
+If your tomograms are already processed and stored locally in .mrc format (e.g., from Warp, IMOD, or AreTomo), you can import them into a new or existing CoPick project using:
+
+```
+octopi import-mrc-volumes \
+    --input-folder /path/to/mrc/files --config /path/to/config.json \
+    --target-tomo-type denoised --input-voxel-size --output-voxel-size 10
+```
+
+octopi also can process tomograms that are hosted on the data portal. Users can download tomograms onto their own remote machine especially if they would like to downsample the tomograms to a lower resolution for speed and memory. You can download and process the tomograms using:
+```
+octopi download-dataportal \
+    --config /path/to/config.json --datasetID 10445 --overlay-path path/to/saved/zarrs \
+    --input-voxel-size 5 --output-voxel-size 10 \
+    --dataportal-name wbp --target-tomotype wbp 
+```
+
+### 📁 Training Labels Preparation  
+
+Use `octopi create-targets` to create semantic masks for proteins of interest using annotation metadata. In this example lets generate picks segmentations for dataset 10439 from the CZ cryoET Dataportal (only need to run this step once). 
+```
+octopi create-targets \
+    --config config.json \
+    --target apoferritin --target beta-galactosidase,slabpick,1 \
+    --target ribosome,pytom,0 --target virus-like-particle,pytom,0 \
+    --seg-target membrane \
+    --tomo-alg wbp --voxel-size 10 \
+    --target-session-id 1 --target-segmentation-name remotetargets \
+    --target-user-id train-octopi
+```
+
+### 🧠 Training a single 3D U-Net model  
+Train a 3D U-Net model on the prepared datasets using the prepared target segmentations. We can use tomograms derived from multiple copick projects. 
+``` 
+octopi train-model \
+    --config experiment,config1.json \
+    --config simulation,config2.json \
+    --voxel-size 10 --tomo-alg wbp --Nclass 8 \
+    --tomo-batch-size 50 --num-epochs 100 --val-interval 10 \
+    --target-info remotetargets,train-octopi,1
+```
+Outputs will include model weights (.pth), logs, and training metrics.
+
+### 🔍 Model exploration with Optuna
+
+octopi🐙 supports automatic neural architecture search using Optuna, enabling efficient discovery of optimal 3D U-Net configurations through Bayesian optimization. This allows users to maximize segmentation accuracy without manual tuning.
+
+To launch a model exploration job:
+```
+octopi model-explore \
+    --config experiment,/mnt/dataportal/ml_challenge/config.json \
+    --config simulation,/mnt/dataportal/synthetic_ml_challenge/config.json \
+    --voxel-size 10 --tomo-alg wbp --Nclass 8 \
+    --model-save-path train_results
+```
+Each trial evaluates a different architecture and logs:
+	•	Segmentation performance metrics
+	•	Model weights and configs
+	•	Training curves and validation loss
+
+🔬 Trials are automatically tracked with MLflow and saved under the specified `--model-save-path`.
+
+#### Optuna Dashboard
+
+To quickly asses the exploration results and observe which trials results the best architectures, Optuna provides a dashboard that summarizes all the information on a dashboard. The instrucutions to access the dashboard are available here - https://optuna-dashboard.readthedocs.io/en/latest/getting-started.html, it is recommended to use either VS-Code extension or CLI. 
+ 
+#### 📊 MLflow experiment tracking   
+
+To use CZI cloud MLflow tracker, add a `.env` in the root directory like below. You can get a CZI MLflow access token from [here](https://mlflow.cw.use4-prod.si.czi.technology/api/2.0/mlflow/users/access-token) (note that a new token will be generated everytime you open this site).
+```
+MLFLOW_TRACKING_USERNAME = <Your_CZ_email>
+MLFLOW_TRACKING_PASSWORD = <Your_mlflow_access_token>
+```
+
+octopi supports MLflow for logging and visualizing model training and hyperparameter search results, including:
+	•	Training loss/validation metrics over time
+	•	Model hyperparameters and architecture details
+	•	Trial comparison (e.g., best performing model)
+
+You can use either a local MLflow instance, a remote (HPC) instance, or the CZI cloud server:
+
+#### 🧪 Local MLflow Dashboard
+
+To inspect results locally: `mlflow ui` and open http://localhost:5000 in your browser.
+
+#### 🖥️ HPC Cluster MLflow Access (Remote via SSH tunnel)
+
+If running octopi on a remote cluster (e.g., Biohub Bruno), forward the MLflow port. 
+On your local machine: 
+ `ssh -L 5000:localhost:5000 remote_username@remote_host` (in the case of Bruno the remote would be `login01.czbiohub.org`). 
+ 
+Then on the remote terminal (login node): ` mlflow ui --host 0.0.0.0 --port 5000` to launch the MLFlow dashboard on a local borwser.
+
+#### ☁️ CZI coreweave cluser
+
+For the CZI coreweave cluser, MLflow is already hosted. Go to the CZI [mlflow server](https://mlflow.cw.use4-prod.si.czi.technology/). 
+
+🔐 A .env file is required to authenticate (see Getting Started section).
+📁 Be sure to register your project name in MLflow before launching runs.
+
+### 🔮 Segmentation
+Generate segmentation prediction masks for tomograms in a given copick project.
+```
+octopi inference \
+    --config config.json \
+    --seg-info predict,unet,1 \
+    --model-config train_results/best_model_config.yaml \
+    --model-weights train_results/best_model.pth \
+    --voxel-size 10 --tomo-alg wbp --tomo-batch-size 25
+```
+Output masks will be saved to the corresponding copick project under the `seg-info` input.
+
+### 📍 Localization
+Convert the segmentation masks into particle coordinates. 
+```
+octopi localize \
+    --config config.json \
+    --pick-session-id 1 --pick-user-id unet \
+    --seg-info predict,unet,1
+```
+
+## Contributing
+
+This project adheres to the Contributor Covenant code of conduct. By participating, you are expected to uphold this code. Please report unacceptable behavior to opensource@chanzuckerberg.com.
+
+## Reporting Security Issues
+
+Please note: If you believe you have found a security issue, please responsibly disclose by contacting us at security@chanzuckerberg.com.
+
+

octopi-1.0/octopi/datasets/augment.py

@@ -0,0 +1,84 @@
+from monai.transforms import (
+    Compose, 
+    RandFlipd, 
+    Orientationd, 
+    RandRotate90d, 
+    NormalizeIntensityd,
+    EnsureChannelFirstd, 
+    RandCropByLabelClassesd,
+    RandScaleIntensityd,
+    RandShiftIntensityd,
+    RandAdjustContrastd,
+    RandGaussianNoised,
+    ScaleIntensityRanged,  
+    RandomOrder,
+)
+
+def get_transforms():
+    """
+    Returns non-random transforms.
+    """
+    return Compose([
+        EnsureChannelFirstd(keys=["image", "label"], channel_dim="no_channel"),
+        NormalizeIntensityd(keys="image"),
+        Orientationd(keys=["image", "label"], axcodes="RAS")
+    ])
+
+def get_random_transforms( input_dim, num_samples, Nclasses):
+    """
+    Input:
+        input_dim: tuple of (nx, ny, nz)
+        num_samples: int
+        Nclasses: int
+
+    Returns random transforms.
+    
+    For data with a missing wedge along the first axis (causing smearing in that direction),
+    we avoid rotations that would move this artifact to other axes. We only rotate around
+    the first axis (spatial_axes=[1, 2]) and avoid flipping along the first axis.
+    """
+    return Compose([
+        RandCropByLabelClassesd(
+            keys=["image", "label"],
+            label_key="label",
+            spatial_size=[input_dim[0], input_dim[1], input_dim[2]],     
+            num_classes=Nclasses,
+            num_samples=num_samples
+        ),
+        # Only rotate around the first axis (keeping the missing wedge orientation consistent)
+        RandRotate90d(keys=["image", "label"], prob=0.5, spatial_axes=[1, 2], max_k=3),
+        # Avoid flipping along the first axis (where the missing wedge is)
+        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),  # Removed
+        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
+        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),        
+        RandomOrder([
+            # Intensity augmentations are still appropriate
+            RandScaleIntensityd(keys="image", prob=0.5, factors=(0.85, 1.15)),
+            RandShiftIntensityd(keys="image", prob=0.5, offsets=(-0.15, 0.15)),
+            RandAdjustContrastd(keys="image", prob=0.5, gamma=(0.85, 1.15)),
+            RandGaussianNoised(keys="image", prob=0.5, mean=0.0, std=0.5),  # Reduced noise std
+        ]),
+    ]) 
+
+    # Augmentations to Explore in the Future: 
+    # Intensity-based augmentations
+    # RandHistogramShiftd(keys="image", prob=0.5, num_control_points=(3, 5))
+    # RandGaussianSmoothd(keys="image", prob=0.5, sigma_x=(0.5, 1.5), sigma_y=(0.5, 1.5), sigma_z=(0.5, 1.5)),
+
+    # Geometric Transforms
+    # RandAffined(
+    #     keys=["image", "label"],
+    #     rotate_range=(0.1, 0.1, 0.1),  # Rotation angles (radians) for x, y, z axes
+    #     scale_range=(0.1, 0.1, 0.1),   # Scale range for isotropic/anisotropic scaling
+    #     prob=0.5,                      # Probability of applying the transform
+    #     padding_mode="border"          # Handle out-of-bounds values
+    # )    
+
+def get_predict_transforms():
+    """
+    Returns predict transforms.
+    """
+    return Compose([
+        EnsureChannelFirstd(keys=["image"], channel_dim="no_channel"),
+        NormalizeIntensityd(keys="image")
+    ])

octopi-1.0/octopi/datasets/cached_datset.py

@@ -0,0 +1,113 @@
+from typing import List, Tuple, Callable, Optional, Dict, Any
+from monai.transforms import Compose
+from monai.data import CacheDataset
+from octopi import io
+from tqdm import tqdm
+import os, sys
+
+class MultiConfigCacheDataset(CacheDataset):
+    """
+    A custom CacheDataset that loads data lazily from multiple sources
+    with consolidated loading and caching process.
+    """
+    
+    def __init__(
+        self,
+        manager,
+        run_ids: List[Tuple[str, str]],
+        transform: Optional[Callable] = None,
+        cache_rate: float = 1.0,
+        num_workers: int = 0,
+        progress: bool = True,
+        copy_cache: bool = True,
+        cache_num: int = sys.maxsize
+    ):
+        # Save reference to manager and run_ids
+        self.manager = manager
+        self.run_ids = run_ids
+        self.progress = progress
+        
+        # Prepare empty data list first - don't load immediately
+        self.data = []
+        
+        # Initialize the parent CacheDataset with an empty list
+        # We'll override the _fill_cache method to handle loading and caching in one step
+        super().__init__(
+            data=[],  # Empty list - we'll load data in _fill_cache
+            transform=transform,
+            cache_rate=cache_rate,
+            num_workers=num_workers,
+            progress=False,  # We'll handle our own progress
+            copy_cache=copy_cache,
+            cache_num=cache_num
+        )
+    
+    def _fill_cache(self):
+        """
+        Override the parent's _fill_cache method to combine loading and caching.
+        """
+        if self.progress:
+            print("Loading and caching dataset...")
+        
+        # Load and process data in a single operation
+        self.data = []
+        iterator = tqdm(self.run_ids, desc="Loading dataset") if self.progress else self.run_ids
+        
+        for session_name, run_name in iterator:
+            root = self.manager.roots[session_name]
+            batch_data = io.load_training_data(
+                root, 
+                [run_name], 
+                self.manager.voxel_size, 
+                self.manager.tomo_algorithm, 
+                self.manager.target_name, 
+                self.manager.target_session_id, 
+                self.manager.target_user_id, 
+                progress_update=False
+            )
+            
+            self.data.extend(batch_data)
+            
+            # Process and cache this batch right away
+            for i, item in enumerate(batch_data):
+                if len(self._cache) < self.cache_num and self.cache_rate > 0.0:
+                    if np.random.random() < self.cache_rate:
+                        self._cache.append(self._transform(item))
+        
+        # Check max label value if needed
+        if hasattr(self.manager, '_check_max_label_value'):
+            self.manager._check_max_label_value(self.data)
+        
+        # Update the _data attribute to match the loaded data
+        self._data = self.data
+    
+    def __len__(self):
+        """
+        Return the length of the dataset.
+        """
+        if not self.data:
+            self._fill_cache()  # Load data if not loaded yet
+        return len(self.data)
+    
+    def __getitem__(self, index):
+        """
+        Return the item at the given index.
+        """
+        if not self.data:
+            self._fill_cache()  # Load data if not loaded yet
+        
+        # Use parent's logic for cached items
+        if index < len(self._cache):
+            return self._cache[index]
+        
+        # Otherwise transform on-the-fly
+        return self._transform(self.data[index])
+
+# TODO: Implement Single Config Cache Dataset
+# class SingleConfigCacheDataset(CacheDataset):
+#     def __init__(self, 
+#                  root: Any, 
+#                  run_ids: List[str], 
+#                  voxel_size: float, 
+#                  tomo_algorithm: str, 
+#                  target_name: str, 

octopi-1.0/octopi/datasets/dataset.py

@@ -0,0 +1,19 @@
+from torch.utils.data import Dataset
+
+class DynamicDataset(Dataset):
+    def __init__(self, data, transform=None):
+        self.data = data
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        sample = self.data[idx]
+        if self.transform:
+            sample = self.transform(sample)
+        return sample
+
+    def update_data(self, new_data):
+        """Update the internal dataset with new data."""
+        self.data = new_data