octopi 1.0__py3-none-any.whl

octopi/submit_slurm.py

@@ -0,0 +1,95 @@
+def create_shellsubmit(
+    job_name, 
+    output_file,
+    shell_name,
+    conda_path,
+    command,
+    num_gpus = 1, 
+    gpu_constraint = 'h100'):
+
+    if num_gpus > 0:
+        slurm_gpus = f'#SBATCH --partition=gpu\n#SBATCH --gpus={gpu_constraint}:{num_gpus}'
+    else:
+        slurm_gpus = f'#SBATCH --partition=cpu'
+
+    shell_script_content = f"""#!/bin/bash
+
+{slurm_gpus}
+#SBATCH --time=18:00:00
+#SBATCH --cpus-per-task=4
+#SBATCH --mem-per-cpu=16G
+#SBATCH --job-name={job_name}
+#SBATCH --output={output_file}
+
+ml anaconda 
+conda activate {conda_path}
+{command}
+"""
+
+    # Save to file
+    with open(shell_name, "w") as file:
+        file.write(shell_script_content)
+
+    print(f"\nShell script has been created successfully as {shell_name}\n")
+
+def create_shellsubmit_array(
+    job_name, 
+    output_file,
+    shell_name,
+    conda_path,
+    command,
+    job_array = [min, max]):
+
+    shell_script_content = f"""#!/bin/bash
+
+#SBATCH --time=18:00:00
+#SBATCH --cpus-per-task=4
+#SBATCH --mem-per-cpu=16G
+#SBATCH --job-name={job_name}
+#SBATCH --output={output_file}
+#SBATCH --array={job_array[0]}-{job_array[1]}
+
+ml anaconda 
+conda activate {conda_path}
+{command}
+"""
+
+    # Save to file
+    with open(shell_name, "w") as file:
+        file.write(shell_script_content)
+
+    print(f"\nShell script has been created successfully as {shell_name}\n")
+
+def create_multiconfig_shellsubmit(
+    job_name, 
+    output_file,
+    shell_name,
+    conda_path,
+    base_inputs,
+    config_inputs,
+    command):
+
+    multiconfig = f"""#! /bin/bash
+
+#SBATCH --job-name={job_name}
+#SBATCH --time=24:00:00
+#SBATCH --cpus-per-task=4
+#SBATCH --mem-per-cpu=12G
+#SBATCH --partition=cpu
+#SBATCH --output={output_file}
+
+ml anaconda 
+{conda_path}
+
+{base_inputs}
+
+{config_inputs}
+
+{command}
+"""
+
+    # Save to file
+    with open(shell_name, "w") as file:
+        file.write(multiconfig)
+
+    print(f"\nShell script has been created successfully as {shell_name}\n")

octopi/utils.py

@@ -0,0 +1,238 @@
+from monai.networks.nets import UNet, AttentionUnet
+from typing import List, Tuple, Union
+from dotenv import load_dotenv
+import argparse, octopi
+import torch, random, os, yaml
+from typing import List
+import numpy as np
+
+##############################################################################################################################
+
+def mlflow_setup():
+
+    module_root = os.path.dirname(octopi.__file__)
+    dotenv_path = module_root.replace('src/octopi','') + '.env'
+    load_dotenv(dotenv_path=dotenv_path)
+
+    # MLflow setup
+    username = os.getenv('MLFLOW_TRACKING_USERNAME')
+    password = os.getenv('MLFLOW_TRACKING_PASSWORD')
+    if not password or not username:
+        print("Password not found in environment, loading from .env file...")
+        load_dotenv()  # Loads environment variables from a .env file
+        username = os.getenv('MLFLOW_TRACKING_USERNAME')
+        password = os.getenv('MLFLOW_TRACKING_PASSWORD')
+    
+    # Check again after loading .env file
+    if not password:
+        raise ValueError("Password is not set in environment variables or .env file!")
+    else:
+        print("Password loaded successfully")
+        os.environ['MLFLOW_TRACKING_USERNAME'] = username
+        os.environ['MLFLOW_TRACKING_PASSWORD'] = password    
+
+    return os.getenv('MLFLOW_TRACKING_URI')
+
+##############################################################################################################################
+
+def set_seed(seed):
+    # Set the seed for Python's random module
+    random.seed(seed)
+
+    # Set the seed for NumPy
+    np.random.seed(seed)
+
+    # Set the seed for PyTorch (both CPU and GPU)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)  # If using multi-GPU
+
+    # Ensure reproducibility of operations by disabling certain optimizations
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    
+###############################################################################################################################
+
+def parse_list(value: str) -> List[str]:
+    """
+    Parse a string representing a list of items.
+    Supports formats like '[item1,item2,item3]' or 'item1,item2,item3'.
+    """
+    value = value.strip("[]")  # Remove brackets if present
+    return [x.strip() for x in value.split(",")]
+
+###############################################################################################################################
+
+def parse_int_list(value: str) -> List[int]:
+    """
+    Parse a string representing a list of integers.
+    Supports formats like '[1,2,3]' or '1,2,3'.
+    """
+    return [int(x) for x in parse_list(value)]
+
+###############################################################################################################################
+
+def string2bool(value: str):
+    """
+    Custom function to convert string values to boolean.
+    """
+    if isinstance(value, bool):
+        return value
+    if value.lower() in {'True', 'true', 't', '1', 'yes'}:
+        return True
+    elif value.lower() in {'False', 'false', 'f', '0', 'no'}:
+        return False
+    else:
+        raise argparse.ArgumentTypeError(f"Invalid boolean value: {value}")
+
+###############################################################################################################################
+
+def parse_target(value: str) -> Tuple[str, Union[str, None], Union[str, None]]:
+    """
+    Parse a single target string.
+    Expected formats:
+      - "name"
+      - "name,user_id,session_id"
+    """
+    parts = value.split(',')
+    if len(parts) == 1:
+        obj_name = parts[0]
+        return obj_name, None, None
+    elif len(parts) == 3:
+        obj_name, user_id, session_id = parts
+        return obj_name, user_id, session_id
+    else:
+        raise argparse.ArgumentTypeError(
+            f"Invalid target format: '{value}'. Expected 'name' or 'name,user_id,session_id'."
+        )
+
+
+def parse_seg_target(value: str) -> List[Tuple[str, Union[str, None], Union[str, None]]]:
+    """
+    Parse segmentation targets. Each target can have the format:
+      - "name"
+      - "name,user_id,session_id"
+    Multiple targets can be comma-separated.
+    """
+    targets = []
+    for target in value.split(';'):  # Use ';' as a separator for multiple targets
+        parts = target.split(',')
+        if len(parts) == 1:
+            name = parts[0]
+            targets.append((name, None, None))
+        elif len(parts) == 3:
+            name, user_id, session_id = parts
+            targets.append((name, user_id, session_id))
+        else:
+            raise argparse.ArgumentTypeError(
+                f"Invalid seg-target format: '{target}'. Expected 'name' or 'name,user_id,session_id'."
+            )
+    return targets
+
+def parse_copick_configs(config_entries: List[str]):
+    """
+    Parse a string representing a list of CoPick configuration file paths.
+    """
+    # Process the --config arguments into a dictionary
+    copick_configs = {}
+
+    for config_entry in config_entries:
+        if ',' in config_entry:
+            # Entry has a session name and a config path
+            try:
+                session_name, config_path = config_entry.split(',', 1)
+                copick_configs[session_name] = config_path
+            except ValueError:
+                raise argparse.ArgumentTypeError(
+                    f"Invalid format for --config entry: '{config_entry}'. Expected 'session_name,/path/to/config.json'."
+                )
+        else:
+            # Single configuration path without a session name
+            # if "default" in copick_configs:
+            #     raise argparse.ArgumentTypeError(
+            #         f"Only one single-path --config entry is allowed when using default configurations. "
+            #         f"Detected duplicate: {config_entry}"
+            #     )
+            # copick_configs["default"] = config_entry
+            copick_configs = config_entry
+
+        # if ',' in config_entry:
+        #     parts = config_entry.split(',')
+        #     if len(parts) == 2:
+        #         # Entry with session name and config path
+        #         session_name, config_path = parts
+        #         copick_configs[session_name] = {"path": config_path, "algorithm": None}
+        #     elif len(parts) == 3:
+        #         # Entry with session name, config path, and algorithm
+        #         session_name, config_path, algorithm = parts
+        #         copick_configs[session_name] = {"path": config_path, "algorithm": algorithm}    
+        # else:
+        #     copick_configs = config_entry
+
+    return copick_configs
+
+def parse_data_split(value: str) -> Tuple[float, float, float]:
+    """
+    Parse data split ratios from string input.
+    
+    Args:
+        value: Either a single float (e.g., "0.8") or two comma-separated floats (e.g., "0.7,0.1")
+    
+    Returns:
+        Tuple of (train_ratio, val_ratio, test_ratio)
+    
+    Examples:
+        "0.8" -> (0.8, 0.2, 0.0)
+        "0.7,0.1" -> (0.7, 0.1, 0.2)
+    """
+    parts = value.split(',')
+    
+    if len(parts) == 1:
+        # Single value provided - use it as train ratio
+        train_ratio = float(parts[0])
+        val_ratio = 1.0 - train_ratio
+        test_ratio = 0.0
+    elif len(parts) == 2:
+        # Two values provided - use as train and val ratios
+        train_ratio = float(parts[0])
+        val_ratio = float(parts[1])
+        test_ratio = 1.0 - train_ratio - val_ratio
+    else:
+        raise ValueError("Data split must be either a single value or two comma-separated values")
+    
+    # Validate ratios
+    if train_ratio < 0 or val_ratio < 0 or test_ratio < 0:
+        raise ValueError("All ratios must be non-negative")
+    
+    if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-6:
+        raise ValueError(f"Ratios must sum to 1.0, got {train_ratio + val_ratio + test_ratio}")
+    
+    return round(train_ratio, 2), round(val_ratio, 2), round(test_ratio, 2)
+
+##############################################################################################################################
+
+# Create a custom dumper that uses flow style for lists only.
+class InlineListDumper(yaml.SafeDumper):
+    def represent_list(self, data):
+        node = super().represent_list(data)
+        node.flow_style = True  # Use inline style for lists
+        return node
+
+def save_parameters_yaml(params: dict, output_path: str):
+    """
+    Save parameters to a YAML file.
+    """
+    InlineListDumper.add_representer(list, InlineListDumper.represent_list)
+    with open(output_path, 'w') as f:
+        yaml.dump(params, f, Dumper=InlineListDumper, default_flow_style=False, sort_keys=False)
+
+def load_yaml(path: str) -> dict:
+    """
+    Load a YAML file and return the contents as a dictionary.
+    """
+    if os.path.exists(path):
+        with open(path, 'r') as f:
+            return yaml.safe_load(f)
+    else:
+        raise FileNotFoundError(f"File not found: {path}")

octopi/visualization_tools.py

@@ -0,0 +1,201 @@
+import matplotlib.colors as mcolors
+from typing import Optional, List
+import matplotlib.pyplot as plt
+from octopi import io
+import numpy as np
+
+# Define the plotting function
+def show_tomo_segmentation(tomo, seg, vol_slice):
+    
+    plt.figure(figsize=(20, 10))
+    
+    # Tomogram
+    plt.subplot(1, 3, 1)
+    plt.title('Tomogram')
+    plt.imshow(tomo[vol_slice], cmap='gray')
+    plt.axis('off')
+    
+    # Painted Segmentation
+    plt.subplot(1, 3, 2)
+    plt.title('Painted Segmentation from Picks')
+    plt.imshow(seg[vol_slice], cmap='viridis')
+    plt.axis('off')
+
+    plt.subplot(1, 3, 3)
+    plt.title('Overlay')
+    plt.imshow(tomo[vol_slice], cmap='gray')
+    plt.imshow(seg[vol_slice], cmap='viridis', alpha=0.5)  # Add alpha=0.5 for 50% transparency
+    plt.axis('off')
+    
+    plt.tight_layout()
+    plt.show()
+
+def show_labeled_tomo_segmentation(tomo, seg, seg_labels, unique_values, vol_slice):
+
+        # # Check unique values in segmentation to ensure correct mapping
+        # unique_values = np.unique(seg)
+
+    plt.figure(figsize=(20, 10))
+    num_classes = len(seg_labels)        
+
+    # Dynamically update the labels and colormap based on unique values
+    seg_labels_filtered = {k: v for k, v in seg_labels.items() if k in unique_values}
+    num_classes = len(seg_labels_filtered)
+
+    # Create a discrete colormap
+    colors = plt.cm.tab20b(np.linspace(0, 1, num_classes))  # You can use other colormaps like 'Set3', 'tab20', etc.
+    cmap = mcolors.ListedColormap(colors)
+    bounds = list(seg_labels_filtered.keys()) + [max(seg_labels_filtered.keys())]
+    # norm = mcolors.BoundaryNorm(bounds, cmap.N)
+
+    # Tomogram plot
+    plt.subplot(1, 2, 1)
+    plt.title('Tomogram')
+    plt.imshow(tomo[vol_slice], cmap='gray')
+    plt.axis('off')
+
+    # Prediction segmentation plot
+    plt.subplot(1, 2, 2)
+    plt.title('Prediction Segmentation')
+    im = plt.imshow(seg[vol_slice], cmap=cmap)  # Use norm and cmap for segmentation
+    plt.axis('off')
+
+    # Add the labeled color bar
+    cbar = plt.colorbar(im, ticks=list(seg_labels_filtered.keys()))
+    cbar.ax.set_yticklabels([seg_labels_filtered[i] for i in seg_labels_filtered.keys()])  # Set custom labels
+
+    plt.tight_layout()
+    plt.show()    
+
+
+def show_tomo_points(tomo, run, objects, user_id, vol_slice, session_id = None, slice_proximity_threshold = 3):
+    plt.figure(figsize=(20, 10))
+
+    plt.imshow(tomo[vol_slice],cmap='gray')
+    plt.axis('off')
+
+    for name,_,_ in objects:
+        try:    
+            coordinates = io.get_copick_coordinates(run, name=name, user_id=user_id, session_id=session_id)
+            close_points = coordinates[np.abs(coordinates[:, 0] - vol_slice) <= slice_proximity_threshold]
+            plt.scatter(close_points[:, 2], close_points[:, 1], label=name, s=15)
+        except:
+            pass
+
+    plt.show()
+
+def compare_tomo_points(tomo, run, objects, vol_slice, user_id1, user_id2, 
+                        session_id1 = None, session_id2 = None, slice_proximity_threshold = 3):
+    plt.figure(figsize=(20, 10))
+
+    plt.subplot(1, 2, 1)
+    plt.imshow(tomo[vol_slice],cmap='gray')
+    plt.title(f'{user_id1} Picks')
+
+    for name,_,_ in objects:
+        try:
+            coordinates = io.get_copick_coordinates(run, name=name, user_id=user_id1, session_id=session_id1)
+            close_points = coordinates[np.abs(coordinates[:, 0] - vol_slice) <= slice_proximity_threshold]
+            plt.scatter(close_points[:, 2], close_points[:, 1], label=name, s=15)    
+        except: 
+            pass
+
+    plt.subplot(1, 2, 2)
+    plt.imshow(tomo[vol_slice],cmap='gray')
+    plt.title(f'{user_id2} Picks')
+    
+    for name,_,_ in objects:
+        try:
+            coordinates = io.get_copick_coordinates(run, name=name, user_id=user_id2, session_id=session_id2)
+            close_points = coordinates[np.abs(coordinates[:, 0] - vol_slice) <= slice_proximity_threshold]
+            plt.scatter(close_points[:, 2], close_points[:, 1], label=name, s=15)
+        except:
+            pass
+
+    plt.axis('off')
+    plt.show()
+
+def plot_training_results(
+    results, 
+    class_names: Optional[List[str]] = None,
+    save_plot: str = None,
+    fig = None, axs = None):
+
+    # Create a 2x2 subplot layout
+    if fig is None:
+        fig, axs = plt.subplots(2, 2, figsize=(12, 10))
+    else:
+        # Clear previos plots
+        for ax in axs.flatten():
+            ax.clear()
+
+    fig.suptitle("Metrics Over Epochs", fontsize=16)
+
+    # Unpack the data for loss (logged every epoch)
+    epochs_loss = [epoch for epoch, _ in results['loss']]
+    loss = [value for _, value in results['loss']]
+    val_epochs_loss = [epoch for epoch, _ in results['val_loss']]
+    val_loss = [value for _,value in results['val_loss']]
+
+    # Plot Training Loss in the top-left
+    axs[0, 0].plot(epochs_loss, loss, label="Training Loss")
+    axs[0, 0].plot(val_epochs_loss, val_loss, label='Validation Loss')
+    axs[0, 0].set_xlabel("Epochs")
+    axs[0, 0].set_ylabel("Loss")
+    axs[0, 0].set_title("Training Loss")
+    axs[0, 0].legend()
+    axs[0, 0].tick_params(axis='both', direction='in', top=True, right=True, length=6, width=1)
+
+    # For metrics that are logged every `val_interval` epochs
+    epochs_metrics = [epoch for epoch, _ in results['avg_recall']]
+    
+    # Determine the number of classes and names
+    num_classes = len([key for key in results.keys() if key.startswith('recall_class')])
+
+    if class_names is None or len(class_names) != num_classes - 1:
+        class_names = [f"Class {i+1}" for i in range(num_classes)]
+
+    # Plot Recall in the top-right
+    for class_idx in range(num_classes):
+        recall_class = [value for _, value in results[f'recall_class{class_idx+1}']]
+        axs[0, 1].plot(epochs_metrics, recall_class, label=f"{class_names[class_idx]}")
+    axs[0, 1].set_xlabel("Epochs")
+    axs[0, 1].set_ylabel("Recall")
+    axs[0, 1].set_title("Recall per Class")
+    # axs[0, 1].legend()
+    axs[0, 1].tick_params(axis='both', direction='in', top=True, right=True, length=6, width=1)
+
+    # Plot Precision in the bottom-left
+    for class_idx in range(num_classes):
+        precision_class = [value for _, value in results[f'precision_class{class_idx+1}']]
+        axs[1, 0].plot(epochs_metrics, precision_class, label=f"{class_names[class_idx]}")
+    axs[1, 0].set_xlabel("Epochs")
+    axs[1, 0].set_ylabel("Precision")
+    axs[1, 0].set_title("Precision per Class")
+    axs[1, 0].legend()
+    axs[1, 0].tick_params(axis='both', direction='in', top=True, right=True, length=6, width=1)
+
+    # Plot F1 Score in the bottom-right
+    for class_idx in range(num_classes):
+        f1_class = [value for _, value in results[f'f1_class{class_idx+1}']]
+        axs[1, 1].plot(epochs_metrics, f1_class, label=f"{class_names[class_idx]}")
+    axs[1, 1].set_xlabel("Epochs")
+    axs[1, 1].set_ylabel("F1 Score")
+    axs[1, 1].set_title("F1 Score per Class")
+    # axs[1, 1].legend()
+    axs[1, 1].tick_params(axis='both', direction='in', top=True, right=True, length=6, width=1)
+
+    # Adjust layout and show plot
+    plt.tight_layout(rect=[0, 0, 1, 0.96])  # Leave space for the main title
+
+    fig.savefig(save_plot)
+    fig.canvas.draw()
+
+    return fig, axs
+
+    # if save_plot: 
+    #     fig.savefig(save_plot)
+    # else:
+    #     plt.show()
+    #     # # Just draw the plot without displaying
+    #     # fig.canvas.draw()

octopi-1.0.dist-info/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.
+```