octopi 1.0__py3-none-any.whl

octopi/entry_points/run_train.py

@@ -0,0 +1,201 @@
+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.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,
+    target_info: Tuple[str, str, str],
+    tomo_algorithm: str = 'wbp',
+    voxel_size: float = 10,
+    trainRunIDs: List[str] = None,
+    validateRunIDs: List[str] = None,    
+    model_config: str = None,
+    model_weights: Optional[str] = None,
+    model_save_path: str = 'results',
+    num_tomo_crops: int = 16,
+    tomo_batch_size: int = 15,
+    lr: float = 1e-3,
+    tversky_alpha: float = 0.5,
+    num_epochs: int = 100,  
+    val_interval: int = 5,
+    best_metric: str = 'avg_f1',
+    data_split: str = '0.8'
+    ):
+
+    # Initialize the data generator to manage training and validation datasets
+    print(f'Training with {copick_config_path}\n')
+    if isinstance(copick_config_path, dict):
+        # Multi-config training
+        data_generator = multi_config_generator.MultiConfigTrainLoaderManager(
+            copick_config_path, 
+            target_info[0], 
+            target_session_id = target_info[2],
+            target_user_id = target_info[1],
+            tomo_algorithm = tomo_algorithm,
+            voxel_size = voxel_size,
+            Nclasses = model_config['num_classes'],
+            tomo_batch_size = tomo_batch_size )
+    else:
+        # Single-config training
+        data_generator = generators.TrainLoaderManager(
+            copick_config_path, 
+            target_info[0], 
+            target_session_id = target_info[2],
+            target_user_id = target_info[1],
+            tomo_algorithm = tomo_algorithm,
+            voxel_size = voxel_size,
+            Nclasses = model_config['num_classes'],
+            tomo_batch_size = tomo_batch_size ) 
+    
+
+    # Get the data splits
+    ratios = utils.parse_data_split(data_split)
+    data_generator.get_data_splits(
+        trainRunIDs = trainRunIDs,
+        validateRunIDs = validateRunIDs,
+        train_ratio = ratios[0], val_ratio = ratios[1], test_ratio = ratios[2],
+        create_test_dataset = False)
+    
+    # Get the reload frequency
+    data_generator.get_reload_frequency(num_epochs)
+
+    # Monai Functions
+    alpha = tversky_alpha
+    beta = 1 - alpha
+    loss_function = TverskyLoss(include_background=True, to_onehot_y=True, softmax=True, alpha=alpha, beta=beta)  
+    metrics_function = ConfusionMatrixMetric(include_background=False, metric_name=["recall",'precision','f1 score'], reduction="none")
+    
+    # Build the Model
+    model_builder = builder.get_model(model_config['architecture'])
+    model = model_builder.build_model(model_config)
+    
+    # Load the Model Weights if Provided 
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if model_weights: 
+        state_dict = torch.load(model_weights, map_location=device, weights_only=True)
+        model.load_state_dict(state_dict)     
+    model.to(device) 
+
+    # Optimizer
+    optimizer = torch.optim.AdamW(model.parameters(), lr, weight_decay=1e-4)
+
+    # Create UNet-Trainer
+    model_trainer = trainer.ModelTrainer(model, device, loss_function, metrics_function, optimizer)
+
+    results = model_trainer.train(
+        data_generator, model_save_path, max_epochs=num_epochs,
+        crop_size=model_config['dim_in'], my_num_samples=num_tomo_crops,
+        val_interval=val_interval, best_metric=best_metric, verbose=True
+    )
+    
+    # Save parameters and results
+    parameters_save_name = os.path.join(model_save_path, "model_config.yaml")
+    io.save_parameters_to_yaml(model_builder, model_trainer, data_generator, parameters_save_name)
+
+    # TODO: Write Results to Zarr or Another File Format? 
+    results_save_name = os.path.join(model_save_path, "results.json")
+    io.save_results_to_json(results, results_save_name)
+
+def train_model_parser(parser_description, add_slurm: bool = False):
+    """
+    Parse the arguments for the training model
+    """
+    parser = argparse.ArgumentParser(
+        description=parser_description,
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    # 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", 
+                             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('--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)")
+    
+    fine_tune_group = parser.add_argument_group("Fine-Tuning Arguments")
+    fine_tune_group.add_argument('--model-config', type=str, help="Path to the model configuration file (typically used for fine-tuning)")
+    fine_tune_group.add_argument('--model-weights', type=str, help="Path to the model weights file (typically used for fine-tuning)") 
+
+    # Model Arguments
+    model_group = parser.add_argument_group("UNet-Model Arguments")
+    common.add_model_parameters(model_group)   
+    
+    # Training Arguments
+    train_group = parser.add_argument_group("Training Arguments")
+    common.add_train_parameters(train_group)
+    
+    # SLURM Arguments
+    if add_slurm:
+        slurm_group = parser.add_argument_group("SLURM Arguments")
+        common.add_slurm_parameters(slurm_group, 'train', gpus = 1)
+
+    args = parser.parse_args()
+    return args
+
+# Entry point with argparse
+def cli():
+    """
+    CLI entry point for training models where results can either be saved to a local directory or a server with MLFlow.
+    """
+
+    # Parse the arguments
+    parser_description = "Train 3D CNN U-Net models"
+    args = train_model_parser(parser_description)
+
+    # Parse the CoPick configuration paths
+    if len(args.config) > 1:    copick_configs = utils.parse_copick_configs(args.config)
+    else:                       copick_configs = args.config[0]
+
+    if args.model_config:
+        model_config = utils.load_yaml(args.model_config)
+    else:
+        model_config = get_model_config(args.channels, args.strides, args.res_units, args.Nclass, args.dim_in)
+
+    # Call the training function
+    train_model(
+        copick_config_path=copick_configs, 
+        target_info=args.target_info,
+        tomo_algorithm=args.tomo_alg,
+        voxel_size=args.voxel_size,
+        model_config=model_config,
+        model_weights=args.model_weights,
+        model_save_path=args.model_save_path,
+        num_tomo_crops=args.num_tomo_crops,
+        tomo_batch_size=args.tomo_batch_size,
+        lr=args.lr,
+        tversky_alpha=args.tversky_alpha,
+        num_epochs=args.num_epochs,
+        val_interval=args.val_interval,
+        best_metric=args.best_metric,
+        trainRunIDs=args.trainRunIDs,
+        validateRunIDs=args.validateRunIDs,
+        data_split=args.data_split
+    )
+
+def get_model_config(channels, strides, res_units, Nclass, dim_in):
+    """
+        Create a model configuration dictionary if no model configuration file is provided.
+    """
+    model_config = {
+        'architecture': 'Unet',
+        'channels': channels,
+        'strides': strides,
+        'num_res_units': res_units, 
+        'num_classes': Nclass,
+        'dropout': 0.1,
+        'dim_in': dim_in
+    }
+    return model_config
+
+if __name__ == "__main__":
+    cli()

octopi/extract/localize.py

@@ -0,0 +1,254 @@
+from skimage.morphology import binary_erosion, binary_dilation, ball
+from scipy.cluster.hierarchy import fcluster, linkage
+from skimage.segmentation import watershed
+from typing import List, Optional, Tuple
+from skimage.measure import regionprops
+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
+
+def processs_localization(run,  
+                          objects, 
+                          seg_info: Tuple[str, str, str],
+                          method: str = 'com', 
+                          voxel_size: float = 10,
+                          filter_size: int = None,
+                          radius_min_scale: float = 0.5, 
+                          radius_max_scale: float = 1.0,
+                          pick_session_id: str = '1',
+                          pick_user_id: str = 'monai'): 
+
+    # Check if method is valid
+    if method not in ['watershed', 'com']:
+        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)
+
+    # Preprocess Segmentation
+    # seg = preprocess_segmentation(seg, voxel_size, objects)
+
+    # If No Segmentation is Found, Return
+    if seg is None:
+        return
+    
+    # Iterate through all user pickable objects
+    for obj in objects:
+
+        # Extract Particle Radius from Root
+        min_radius = obj[2] * radius_min_scale / voxel_size
+        max_radius = obj[2] * radius_max_scale / voxel_size
+
+        if method == 'watershed':
+            points = extract_particle_centroids_via_watershed(seg, obj[1], filter_size, min_radius, max_radius)
+        elif method == 'com': 
+            points = extract_particle_centroids_via_com(seg, obj[1], min_radius, max_radius)
+        points = np.array(points)
+
+        # Save Coordinates if any 3D points are provided
+        if points.size > 2:
+
+            # Remove Picks that are too close to each other
+            # points = remove_repeated_picks(points, min_radius, pixelSize = voxel_size)
+
+            # Swap the coordinates to match the expected format
+            points = points[:,[2,1,0]] 
+
+            # Convert the Picks back to Angstrom
+            points *= voxel_size
+
+            # Save Picks
+            try:
+                picks = run.new_picks(object_name = obj[0], session_id = pick_session_id, user_id=pick_user_id)
+            except:
+                picks = run.get_picks(object_name = obj[0], session_id = pick_session_id, user_id=pick_user_id)[0]
+
+            # Assign Identity As Orientation
+            orientations = np.zeros([points.shape[0], 4, 4])
+            orientations[:,:3,:3] = np.identity(3)
+            orientations[:,3,3] = 1
+
+            picks.from_numpy( points, orientations )
+        else:
+            print(f"{run.name} didn't have any available picks for {obj[0]}!")
+
+
+def extract_particle_centroids_via_watershed(
+        segmentation, 
+        segmentation_idx, 
+        maxima_filter_size, 
+        min_particle_radius, 
+        max_particle_radius):
+    """
+    Process a specific label in the segmentation, extract centroids, and save them as picks.
+
+    Args:
+        segmentation (np.ndarray): Multilabel segmentation array.
+        segmentation_idx (int): The specific label from the segmentation to process.
+        maxima_filter_size (int): Size of the maximum detection filter.
+        min_particle_size (int): Minimum size threshold for particles.
+        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!')
+
+    # 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)
+
+    # Skip if the segmentation label is not present
+    if np.sum(binary_mask) == 0:
+        print(f"No segmentation with label {segmentation_idx} found.")
+        return
+
+    # 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
+    distance = ndi.distance_transform_edt(dilated)
+    local_max = (distance == ndi.maximum_filter(distance, footprint=np.ones((maxima_filter_size, maxima_filter_size, maxima_filter_size))))
+
+    # Watershed segmentation
+    markers, _ = ndi.label(local_max)
+    watershed_labels = watershed(-distance, markers, mask=dilated)
+
+    # Extract region properties and filter based on particle size
+    all_centroids = []
+    for region in regionprops(watershed_labels):
+        if min_particle_size <= region.area <= max_particle_size:
+
+            # Option 1: Use all centroids
+            all_centroids.append(region.centroid)
+
+    return all_centroids
+
+def extract_particle_centroids_via_com(
+        segmentation, 
+        segmentation_idx, 
+        min_particle_radius, 
+        max_particle_radius
+    ):
+    """
+    Process a specific label in the segmentation, extract centroids, and save them as picks.
+
+    Args:
+        segmentation (np.ndarray): Multilabel segmentation array.
+        segmentation_idx (int): The specific label from the segmentation to process.
+        min_particle_size (int): Minimum size threshold for particles.
+        max_particle_size (int): Maximum size threshold for particles.
+    """
+
+    # 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
+    label_objs, _ = ndi.label(segmentation == segmentation_idx)
+
+    # Filter Candidates based on Object Size
+    # Get the sizes of all objects
+    object_sizes = np.bincount(label_objs.flat)
+
+    # Filter the objects based on size
+    valid_objects = np.where((object_sizes > min_particle_size) & (object_sizes < max_particle_size))[0]                        
+
+    # Estimate Coordiantes from CoM for LabelMaps
+    octopiCoords = []
+    for object_num in tqdm(valid_objects):
+        com = ndi.center_of_mass(label_objs == object_num)
+        swapped_com = (com[2], com[1], com[0])
+        octopiCoords.append(swapped_com)
+   
+    return octopiCoords
+
+def remove_repeated_picks(coordinates, distanceThreshold, pixelSize = 1):
+
+    # Calculate the distance matrix for the 3D coordinates
+    dist_matrix = distance.cdist(coordinates[:, :3]/pixelSize, coordinates[:, :3]/pixelSize)
+
+    # Create a linkage matrix using single linkage method
+    Z = linkage(dist_matrix, method='complete')
+
+    # Form flat clusters with a distance threshold to determine groups
+    clusters = fcluster(Z, t=distanceThreshold, criterion='distance')
+
+    # Initialize an array to store the average of each group
+    unique_coordinates = np.zeros((max(clusters), coordinates.shape[1]))
+
+    # Calculate the mean for each cluster
+    for i in range(1, max(clusters) + 1):
+        unique_coordinates[i-1] = np.mean(coordinates[clusters == i], axis=0)
+
+    return unique_coordinates
+
+def preprocess_segmentation(segmentation, voxel_size, particle_info):
+    """
+    Remove tiny fragments that aren't real particles
+    
+    Args:
+        segmentation (np.ndarray): The multilabel segmentation array
+        particle_info (list): List of tuples containing (name, segment_id, radius)
+        
+    Returns:
+        np.ndarray: Processed segmentation with small fragments removed
+    """
+    import numpy as np
+    from skimage.morphology import remove_small_objects
+    
+    processed_seg = segmentation.copy()
+    
+    # Map segment IDs to particle types and their minimum sizes
+    segment_to_info = {}
+    for name, segment_id, radius in particle_info:
+        # # For small particles, use a larger minimum size
+        # if radius < 135:
+        #     scale = 0.65
+        # # Normal threshold for other particles
+        # else:
+        #     scale = 0.4
+        scale = 0.3
+        radius = radius / voxel_size
+        min_size = (4/3) * np.pi * ((radius * 0.5) ** 3)
+        
+        segment_to_info[segment_id] = {
+            'name': name,
+            'min_size': min_size
+        }
+    
+    # Get unique labels
+    unique_labels = np.unique(segmentation)
+    unique_labels = unique_labels[unique_labels > 0]  # Skip background
+    
+    # Process each label
+    for label in unique_labels:
+        if label not in segment_to_info:
+            continue
+            
+        # Create binary mask for this label
+        mask = segmentation == label
+        
+        # Get minimum size for this particle type
+        min_size = segment_to_info[label]['min_size']
+        
+        # Remove small objects
+        cleaned_mask = remove_small_objects(mask, min_size=min_size * scale)
+        
+        # Update segmentation
+        processed_seg[mask & ~cleaned_mask] = 0
+    
+    return processed_seg