octopi 1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

octopi/entry_points/run_optuna.py

@@ -1,7 +1,7 @@
 from octopi.pytorch.model_search_submitter import ModelSearchSubmit
 from octopi.entry_points import common
+from octopi.utils import parsers, io
 import argparse, os, pprint
-from octopi import utils
 
 def optuna_parser(parser_description, add_slurm: bool = False):
     """
@@ -20,22 +20,22 @@ def optuna_parser(parser_description, add_slurm: bool = False):
     # Input Arguments
     input_group = parser.add_argument_group("Input Arguments")
     common.add_config(input_group, single_config=False)
-    input_group.add_argument("--target-info", type=utils.parse_target, default="targets,octopi,1", 
+    input_group.add_argument("--target-info", type=parsers.parse_target, default="targets,octopi,1", 
                              help="Target information, e.g., 'name' or 'name,user_id,session_id'")    
     input_group.add_argument("--tomo-alg", default='wbp', 
                              help="Tomogram algorithm used for training")
     input_group.add_argument("--mlflow-experiment-name", type=str, default="model-search", required=False, 
                              help="Name of the MLflow experiment (default: 'model-search').")
-    input_group.add_argument("--trainRunIDs", type=utils.parse_list, default=None, required=False, 
+    input_group.add_argument("--trainRunIDs", type=parsers.parse_list, default=None, required=False, 
                              help="List of training run IDs, e.g., run1,run2 or [run1,run2].")
-    input_group.add_argument("--validateRunIDs", type=utils.parse_list, default=None, required=False, 
+    input_group.add_argument("--validateRunIDs", type=parsers.parse_list, default=None, required=False, 
                              help="List of validation run IDs, e.g., run3,run4 or [run3,run4].")    
     input_group.add_argument('--data-split', type=str, default='0.8', help="Data split ratios. Either a single value (e.g., '0.8' for 80/20/0 split) "
                                 "or two comma-separated values (e.g., '0.7,0.1' for 70/10/20 split)")
 
     model_group = parser.add_argument_group("Model Arguments")
     model_group.add_argument("--model-type", type=str, default='Unet', required=False, 
-                             choices=['Unet', 'AttentionUnet'],
+                             choices=['Unet', 'AttentionUnet', 'MedNeXt', 'SegResNet'],
                              help="Model type to use for training")
     model_group.add_argument("--Nclass", type=int, default=3, required=False, help="Number of prediction classes in the model")
 
@@ -61,7 +61,7 @@ def cli():
     args = optuna_parser(description)
 
     # Parse the CoPick configuration paths
-    if len(args.config) > 1:    copick_configs = utils.parse_copick_configs(args.config)
+    if len(args.config) > 1:    copick_configs = parsers.parse_copick_configs(args.config)
     else:                       copick_configs = args.config[0]
 
     # Create the model exploration directory
@@ -133,7 +133,7 @@ def save_parameters(args: argparse.Namespace,
     pprint.pprint(params); print()
 
     # Save to YAML file
-    utils.save_parameters_yaml(params, output_path)
+    io.save_parameters_yaml(params, output_path)
 
 if __name__ == "__main__":
     cli()

octopi/entry_points/run_segment_predict.py

@@ -1,8 +1,8 @@
+import torch, argparse, json, pprint, yaml, os
 from octopi.pytorch import segmentation
 from octopi.entry_points import common
-import torch, argparse, json, pprint, yaml, os
-from octopi import utils
 from typing import List, Tuple
+from octopi.utils import io
 
 def inference(
     copick_config_path: str,
@@ -136,7 +136,7 @@ def save_parameters(args: argparse.Namespace,
                     output_path: str):  
 
     # Load the model config
-    model_config = utils.load_yaml(args.model_config)
+    model_config = io.load_yaml(args.model_config)
 
     # Create parameters dictionary
     params = {
@@ -160,7 +160,7 @@ def save_parameters(args: argparse.Namespace,
     pprint.pprint(params); print()
 
     # Save to YAML file
-    utils.save_parameters_yaml(params, output_path)
+    io.save_parameters_yaml(params, output_path)
 
 if __name__ == "__main__":
     cli()

octopi/entry_points/run_train.py

@@ -2,12 +2,11 @@ from octopi.datasets import generators, multi_config_generator
 from monai.losses import DiceLoss, FocalLoss, TverskyLoss
 from octopi.models import common as builder
 from monai.metrics import ConfusionMatrixMetric
+from octopi.utils import parsers, io
 from octopi.entry_points import common 
 from octopi.pytorch import trainer 
-from octopi import io, utils
 import torch, os, argparse
 from typing import List, Optional, Tuple
-import pprint
 
 def train_model(
     copick_config_path: str,
@@ -56,7 +55,7 @@ def train_model(
     
 
     # Get the data splits
-    ratios = utils.parse_data_split(data_split)
+    ratios = parsers.parse_data_split(data_split)
     data_generator.get_data_splits(
         trainRunIDs = trainRunIDs,
         validateRunIDs = validateRunIDs,
@@ -114,11 +113,11 @@ def train_model_parser(parser_description, add_slurm: bool = False):
     # Input Arguments
     input_group = parser.add_argument_group("Input Arguments")
     common.add_config(input_group, single_config=False)
-    input_group.add_argument("--target-info", type=utils.parse_target, default="targets,octopi,1", 
+    input_group.add_argument("--target-info", type=parsers.parse_target, default="targets,octopi,1", 
                              help="Target information, e.g., 'name' or 'name,user_id,session_id'. Default is 'targets,octopi,1'.")
     input_group.add_argument("--tomo-alg", default='wbp', help="Tomogram algorithm used for training")
-    input_group.add_argument("--trainRunIDs", type=utils.parse_list, help="List of training run IDs, e.g., run1,run2,run3")
-    input_group.add_argument("--validateRunIDs", type=utils.parse_list, help="List of validation run IDs, e.g., run4,run5,run6")
+    input_group.add_argument("--trainRunIDs", type=parsers.parse_list, help="List of training run IDs, e.g., run1,run2,run3")
+    input_group.add_argument("--validateRunIDs", type=parsers.parse_list, help="List of validation run IDs, e.g., run4,run5,run6")
     input_group.add_argument('--data-split', type=str, default='0.8', help="Data split ratios. Either a single value (e.g., '0.8' for 80/20/0 split) "
                                 "or two comma-separated values (e.g., '0.7,0.1' for 70/10/20 split)")
     
@@ -153,11 +152,11 @@ def cli():
     args = train_model_parser(parser_description)
 
     # Parse the CoPick configuration paths
-    if len(args.config) > 1:    copick_configs = utils.parse_copick_configs(args.config)
+    if len(args.config) > 1:    copick_configs = parsers.parse_copick_configs(args.config)
     else:                       copick_configs = args.config[0]
 
     if args.model_config:
-        model_config = utils.load_yaml(args.model_config)
+        model_config = io.load_yaml(args.model_config)
     else:
         model_config = get_model_config(args.channels, args.strides, args.res_units, args.Nclass, args.dim_in)
 

octopi/extract/localize.py

@@ -3,15 +3,15 @@ from scipy.cluster.hierarchy import fcluster, linkage
 from skimage.segmentation import watershed
 from typing import List, Optional, Tuple
 from skimage.measure import regionprops
+from copick_utils.io import readers
 from scipy.spatial import distance
 from dataclasses import dataclass
-from octopi import io
 import scipy.ndimage as ndi
 from tqdm import tqdm
 import numpy as np
-import math
+import gc
 
-def processs_localization(run,  
+def process_localization(run,  
                           objects, 
                           seg_info: Tuple[str, str, str],
                           method: str = 'com', 
@@ -27,12 +27,12 @@ def processs_localization(run,
         raise ValueError(f"Invalid method '{method}'. Expected 'watershed' or 'com'.")
 
     # Get Segmentation
-    seg = io.get_segmentation_array(run, 
-                                    voxel_size, 
-                                    seg_info[0], 
-                                    user_id=seg_info[1], 
-                                    session_id=seg_info[2],
-                                    raise_error=False)
+    seg = readers.segmentation(
+        run, voxel_size, 
+        seg_info[0], 
+        user_id=seg_info[1], 
+        session_id=seg_info[2],
+        raise_error=False)
 
     # Preprocess Segmentation
     # seg = preprocess_segmentation(seg, voxel_size, objects)
@@ -99,15 +99,15 @@ def extract_particle_centroids_via_watershed(
         max_particle_size (int): Maximum size threshold for particles.
     """
 
-    if maxima_filter_size is None or maxima_filter_size < 0:
-        AssertionError('Enter a Non-Zero Filter Size!')
+    if maxima_filter_size is None or maxima_filter_size <= 0:
+        raise ValueError('Enter a Non-Zero Filter Size!')
 
     # Calculate minimum and maximum particle volumes based on the given radii
     min_particle_size = (4 / 3) * np.pi * (min_particle_radius ** 3) 
     max_particle_size = (4 / 3) * np.pi * (max_particle_radius ** 3)
 
     # Create a binary mask for the specific segmentation label
-    binary_mask = (segmentation == segmentation_idx).astype(int)
+    binary_mask = (segmentation == segmentation_idx).astype(np.uint8)
 
     # Skip if the segmentation label is not present
     if np.sum(binary_mask) == 0:
@@ -117,6 +117,7 @@ def extract_particle_centroids_via_watershed(
     # Structuring element for erosion and dilation
     struct_elem = ball(1)
     eroded = binary_erosion(binary_mask, struct_elem)
+
     dilated = binary_dilation(eroded, struct_elem)
 
     # Distance transform and local maxima detection
@@ -125,7 +126,13 @@ def extract_particle_centroids_via_watershed(
 
     # Watershed segmentation
     markers, _ = ndi.label(local_max)
+    del local_max
+    gc.collect()
+
     watershed_labels = watershed(-distance, markers, mask=dilated)
+    distance, markers, dilated = None, None, None
+    del distance, markers, dilated
+    gc.collect()
 
     # Extract region properties and filter based on particle size
     all_centroids = []

octopi/extract/membranebound_extract.py

@@ -1,5 +1,5 @@
 from scipy.spatial.transform import Rotation as R
-from octopi import utils, io
+from copick_utils.io import readers
 import scipy.ndimage as ndi
 from typing import Tuple
 import numpy as np
@@ -36,7 +36,7 @@ def process_membrane_bound_extract(run,
     new_session_id = str(int(save_session_id) + 1)  # Convert to string after increment                                  
 
     # Need Better Error Handing for Missing Picks
-    coordinates = io.get_copick_coordinates(
+    coordinates = readers.coordinates(
         run, 
         picks_info[0], picks_info[1], picks_info[2],
         voxel_size,
@@ -54,12 +54,13 @@ def process_membrane_bound_extract(run,
     if membrane_info is None:
         # Flag to distinguish between organelle and membrane segmentation
         membranes_provided = False
-        seg = io.get_segmentation_array(run, 
-                                    voxel_size, 
-                                    organelle_info[0],
-                                    user_id=organelle_info[1], 
-                                    session_id=organelle_info[2],
-                                    raise_error=False)
+        seg = readers.segmentation(
+            run, 
+            voxel_size, 
+            organelle_info[0],
+            user_id=organelle_info[1], 
+            session_id=organelle_info[2],
+            raise_error=False)
         # If No Segmentation is Found, Return
         if seg is None: return     
         elif nPoints == 0 or np.unique(seg).max() == 0:
@@ -68,7 +69,7 @@ def process_membrane_bound_extract(run,
     else:
         # Read both Organelle and Membrane Segmentations
         membranes_provided = True
-        seg = io.get_segmentation_array(
+        seg = readers.segmentation(
             run, 
             voxel_size, 
             membrane_info[0],
@@ -76,7 +77,7 @@ def process_membrane_bound_extract(run,
             session_id=membrane_info[2],
             raise_error=False)
 
-        organelle_seg = io.get_segmentation_array(
+        organelle_seg = readers.segmentation(
             run, 
             voxel_size, 
             organelle_info[0],

octopi/extract/midpoint_extract.py

@@ -1,6 +1,6 @@
 from octopi.extract import membranebound_extract as extract
 from scipy.spatial.transform import Rotation as R
-from octopi import io
+from copick_utils.io import readers
 from scipy.spatial import cKDTree
 from typing import Tuple
 import numpy as np
@@ -28,7 +28,7 @@ def process_midpoint_extract(
     """
 
     # Pull Picks that Are used for Midpoint Extraction
-    coordinates = io.get_copick_coordinates(
+    coordinates = readers.coordinates(
         run, 
         picks_info[0], picks_info[1], picks_info[2],
         voxel_size
@@ -40,7 +40,7 @@ def process_midpoint_extract(
     save_picks_info[2] = save_session_id
 
     # Get Organelle Segmentation
-    seg = io.get_segmentation_array(
+    seg = readers.segmentation(
         run, 
         voxel_size, 
         organelle_info[0],

octopi/main.py

@@ -33,7 +33,7 @@ def cli_main():
         "create-targets": (create_targets, "Generate segmentation targets from coordinates."),
         "train": (train_model, "Train a single U-Net model."),
         "model-explore": (model_explore, "Explore model architectures with Optuna / Bayesian Optimization."),
-        "inference": (inference, "Perform segmentation inference on tomograms."),
+        "segment": (inference, "Perform segmentation inference on tomograms."),
         "localize": (localize, "Perform localization of particles in tomograms."),
         "extract-mb-picks": (extract_mb_picks, "Extract MB Picks from tomograms."),
         "evaluate": (evaluate, "Evaluate the performance of a model."),

octopi/models/common.py

@@ -1,5 +1,5 @@
 from monai.losses import FocalLoss, TverskyLoss
-from octopi import losses
+from octopi.utils import losses
 from octopi.models import (
     Unet, AttentionUnet, MedNeXt, SegResNet
 )

octopi/processing/create_targets_from_picks.py

@@ -1,6 +1,5 @@
 from octopi.processing.segmentation_from_picks import from_picks
-import octopi.processing.writers as write
-from octopi import io 
+from copick_utils.io import readers, writers
 from typing import List
 from tqdm import tqdm
 import numpy as np
@@ -42,7 +41,11 @@ def generate_targets(
 
     # If runIDs are not provided, load all runs
     if run_ids is None:
-        run_ids = [run.name for run in root.runs]
+        run_ids = [run.name for run in root.runs if run.get_voxel_spacing(voxel_size) is not None]
+        skipped_run_ids = [run.name for run in root.runs if run.get_voxel_spacing(voxel_size) is None]
+        
+        if skipped_run_ids:
+            print(f"Warning: skipping runs with no voxel spacing {voxel_size}: {skipped_run_ids}")
 
     # Iterate Over All Runs
     for runID in tqdm(run_ids):
@@ -52,7 +55,7 @@ def generate_targets(
         run = root.get_run(runID)
 
         # Get Tomogram 
-        tomo = io.get_tomogram_array(run, voxel_size, tomo_algorithm)
+        tomo = readers.tomogram(run, voxel_size, tomo_algorithm)
         
         # Initialize Target Volume
         target = np.zeros(tomo.shape, dtype=np.uint8)
@@ -87,6 +90,9 @@ def generate_targets(
                     session_id=train_targets[target_name]["session_id"],
                 )
 
+        # Filter out empty picks
+        query = [pick for pick in query if pick.points is not None]
+
         # Add Picks to Target  
         for pick in query:
             numPicks += len(pick.points)
@@ -100,7 +106,7 @@ def generate_targets(
         # Write Segmentation for non-empty targets
         if target.max() > 0 and numPicks > 0:
             tqdm.write(f'Annotating {numPicks} picks in {runID}...')    
-            write.segmentation(run, target, target_user_name, 
+            writers.segmentation(run, target, target_user_name, 
                                name = target_segmentation_name, session_id= target_session_id, 
                                voxel_size = voxel_size)
     print('Creation of targets complete!')

octopi/processing/downsample.py

@@ -102,11 +102,6 @@ class FourierRescale:
         """
         in_depth, in_height, in_width = volume.shape[-3:]
         
-        # Check if dimensions are odd
-        d_is_odd = in_depth % 2
-        h_is_odd = in_height % 2
-        w_is_odd = in_width % 2
-        
         # Calculate new dimensions
         extent_depth = in_depth * self.input_voxel_size[0]
         extent_height = in_height * self.input_voxel_size[1]
@@ -121,9 +116,10 @@ class FourierRescale:
         new_height = new_height - (new_height % 2)
         new_width = new_width - (new_width % 2)
         
-        # Calculate starting points with odd/even correction
-        start_d = (in_depth - new_depth) // 2 + (d_is_odd)
-        start_h = (in_height - new_height) // 2 + (h_is_odd)
-        start_w = (in_width - new_width) // 2 + (w_is_odd)
+        # Calculate starting points - properly centered around DC component
+        # No odd/even correction needed - just center the crop
+        start_d = (in_depth - new_depth) // 2
+        start_h = (in_height - new_height) // 2
+        start_w = (in_width - new_width) // 2
         
-        return start_d, start_h, start_w, new_depth, new_height, new_width
+        return start_d, start_h, start_w, new_depth, new_height, new_width

octopi/processing/evaluate.py

@@ -1,7 +1,7 @@
-from octopi import utils, io
+from copick_utils.io import readers
 from scipy.spatial import distance
+import copick, json, os, yaml   
 from typing import List
-import copick, json, os
 import numpy as np
 
 class evaluator:
@@ -95,12 +95,12 @@ class evaluator:
             for name, radius in self.objects:
                 
                 # Get Ground Truth and Predicted Coordinates
-                gt_coordinates = io.get_copick_coordinates(
+                gt_coordinates = readers.coordinates(
                     run, name, 
                     self.ground_truth_user_id, self.ground_truth_session_id, 
                     self.voxel_size, raise_error=False
                 )
-                pred_coordinates = io.get_copick_coordinates(
+                pred_coordinates = readers.coordinates(
                     run, name,
                     self.prediction_user_id, self.predict_session_id, 
                     self.voxel_size, raise_error=False
@@ -202,14 +202,27 @@ class evaluator:
             }    
             
             os.makedirs(save_path, exist_ok=True)
-            summary_metrics = { "input": self.input_params, "parameters": self.parameters, 
-                                    "summary_metrics": final_summary_metrics }
-            with open(os.path.join(save_path, 'average_metrics.json'), 'w') as f:
-                json.dump(summary_metrics, f, indent=4)
-            print(f'\nAverage Metrics saved to {os.path.join(save_path, "average_metrics.json")}')
+            summary_metrics = { "input": self.input_params, 
+                                "final_fbeta_score": final_fbeta,  
+                                "aggregated_particle_scores": {    # Optionally add per-particle details
+                                    name: {
+                                        "tp": counts['total_tp'],
+                                        "fp": counts['total_fp'], 
+                                        "fn": counts['total_fn'],
+                                        "weight": self.weights.get(name, 1)
+                                    } for name, counts in aggregated_counts.items()
+                                },
+                                "summary_metrics": final_summary_metrics, 
+                                "parameters": self.parameters,  }
+
+            # Save average metrics to YAML file
+            with open(os.path.join(save_path, 'average_metrics.yaml'), 'w') as f:
+                yaml.dump(summary_metrics, f, indent=4, default_flow_style=False, sort_keys=False)
+            print(f'\nAverage Metrics saved to {os.path.join(save_path, "average_metrics.yaml")}')
             
-            detailed_metrics = { "input": self.input_params, "parameters": self.parameters, 
-                                "metrics": metrics }
+            detailed_metrics = { "input": self.input_params,  
+                                  "metrics": metrics,
+                                 "parameters": self.parameters, }
             with open(os.path.join(save_path, 'metrics.json'), 'w') as f:
                 json.dump(detailed_metrics, f, indent=4)
             print(f'Metrics saved to {os.path.join(save_path, "metrics.json")}')

octopi/processing/importers.py

@@ -1,7 +1,7 @@
 from octopi.processing.downsample import FourierRescale
 import copick, argparse, mrcfile, glob, os
-import octopi.processing.writers as write
 from octopi.entry_points import common
+from copick_utils.io import writers
 from tqdm import tqdm
 
 def from_dataportal(
@@ -57,10 +57,10 @@ def from_dataportal(
 
             # If we want to save the tomograms at a different voxel size, we need to rescale the tomograms
             if output_voxel_size is None:
-                write.tomogram(run, vol, input_voxel_size, target_tomo_type)
+                writers.tomogram(run, vol, input_voxel_size, target_tomo_type)
             else:
                 vol = rescale.run(vol)
-                write.tomogram(run, vol, output_voxel_size, target_tomo_type)
+                writers.tomogram(run, vol, output_voxel_size, target_tomo_type)
     
     print(f'Downloading Complete!! Downloaded {len(root.runs)} runs')
 
@@ -168,7 +168,7 @@ def from_mrcs(
             voxel_size_to_write = input_voxel_size
 
         # Write the tomogram
-        write.tomogram(run, vol, voxel_size_to_write, target_tomo_type)
+        writers.tomogram(run, vol, voxel_size_to_write, target_tomo_type)
     print(f"Processed {len(mrc_files)} files from {mrcs_path}")
 
 

octopi/pytorch/hyper_search.py

@@ -1,10 +1,9 @@
-from monai.losses import FocalLoss, TverskyLoss
 from monai.metrics import ConfusionMatrixMetric
 from octopi.pytorch import trainer
 from mlflow.tracking import MlflowClient
 from octopi.models import common
-from octopi import io, losses
 import torch, mlflow, optuna, gc
+from octopi.utils import io
 
 class BayesianModelSearch:
 
@@ -207,7 +206,7 @@ class BayesianModelSearch:
         if score > best_score_so_far:
             torch.save(model_trainer.model_weights, f'{self.results_dir}/best_model.pth')
             io.save_parameters_to_yaml(self.model_builder, model_trainer, self.data_generator, 
-                                    f'{self.results_dir}/best_model_config.yaml')
+                                    f'{self.results_dir}/model_config.yaml')
 
     def get_best_score(self, trial):
         """Retrieve the best score from the trial."""

octopi/pytorch/model_search_submitter.py

@@ -1,7 +1,7 @@
 from octopi.datasets import generators, multi_config_generator
+from octopi.utils import config, parsers
 from octopi.pytorch import hyper_search
 import torch, mlflow, optuna
-from octopi import utils
 from typing import List
 import pandas as pd
 
@@ -16,16 +16,16 @@ class ModelSearchSubmit:
         voxel_size: float,
         Nclass: int,
         model_type: str,
-        mlflow_experiment_name: str,
-        random_seed: int,
-        num_epochs: int,
-        num_trials: int,
-        tomo_batch_size: int,
-        best_metric: str,
-        val_interval: int,
-        trainRunIDs: List[str],
-        validateRunIDs: List[str],
-        data_split: str
+        best_metric: str = 'avg_f1',
+        num_epochs: int = 1000,
+        num_trials: int = 100,
+        data_split: str = 0.8,
+        random_seed: int = 42,
+        val_interval: int = 10,
+        tomo_batch_size: int = 15,
+        trainRunIDs: List[str] = None,
+        validateRunIDs: List[str] = None,
+        mlflow_experiment_name: str = 'explore',
     ):
         """
         Initialize the ModelSearch class for architecture search with Optuna.
@@ -75,7 +75,7 @@ class ModelSearchSubmit:
         self.data_generator = None
 
         # Set random seed for reproducibility
-        utils.set_seed(self.random_seed)
+        config.set_seed(self.random_seed)
 
         # Initialize dataset generator
         self._initialize_data_generator()
@@ -108,7 +108,7 @@ class ModelSearchSubmit:
             )
 
         # Split datasets into training and validation
-        ratios = utils.parse_data_split(self.data_split)
+        ratios = parsers.parse_data_split(self.data_split)
         self.data_generator.get_data_splits(
             trainRunIDs=self.trainRunIDs,
             validateRunIDs=self.validateRunIDs,
@@ -134,7 +134,7 @@ class ModelSearchSubmit:
 
         # Set up MLflow tracking
         try:
-            tracking_uri = utils.mlflow_setup()
+            tracking_uri = config.mlflow_setup()
             mlflow.set_tracking_uri(tracking_uri)
         except Exception as e:
             print(f'Failed to set up MLflow tracking: {e}')
@@ -207,7 +207,7 @@ class ModelSearchSubmit:
             # Run multi-GPU optimization
             study = self.get_optuna_study()
             study.optimize(
-                lambda trial: BayesianModelSearch(self.data_generator, self.model_type).multi_gpu_objective(
+                lambda trial: hyper_search.BayesianModelSearch(self.data_generator, self.model_type).multi_gpu_objective(
                     parent_run, trial,
                     self.num_epochs,
                     best_metric=self.best_metric,