octopi 1.0__py3-none-any.whl

octopi/processing/downsample.py

@@ -0,0 +1,129 @@
+import numpy as np
+import torch
+
+class FourierRescale:
+    def __init__(self, input_voxel_size, output_voxel_size):
+        """
+        Initialize the FourierRescale operation with voxel sizes.
+
+        Parameters:
+            input_voxel_size (int or tuple): Physical spacing of the input voxels (d, h, w)
+                                             or a single int (which will be applied to all dimensions).
+            output_voxel_size (int or tuple): Desired physical spacing of the output voxels (d, h, w)
+                                              or a single int (which will be applied to all dimensions).
+                                              Must be greater than or equal to input_voxel_size.
+        """
+        # Convert to tuples if single int is provided.
+        if isinstance(input_voxel_size, int) or isinstance(input_voxel_size, float):
+            input_voxel_size = (input_voxel_size, input_voxel_size, input_voxel_size)
+        if isinstance(output_voxel_size, int) or isinstance(output_voxel_size, float):
+            output_voxel_size = (output_voxel_size, output_voxel_size, output_voxel_size)
+            
+        self.input_voxel_size = input_voxel_size
+        self.output_voxel_size = output_voxel_size
+        
+        # Check: output voxel size must be greater than or equal to input voxel size (element-wise).
+        if any(out_vs < in_vs for in_vs, out_vs in zip(input_voxel_size, output_voxel_size)):
+            raise ValueError("Output voxel size must be greater than or equal to the input voxel size.")
+
+        # Determine device: use GPU if available, otherwise CPU.
+        if torch.cuda.is_available():
+            self.device = torch.device('cuda')
+        else:
+            self.device = torch.device('cpu')
+
+    def run(self, volume):
+        """
+        Rescale a 3D volume (or a batch of volumes on GPU) using Fourier cropping.
+        """
+        # Initialize return_numpy flag
+        return_numpy = False
+
+        # If a numpy array is passed, convert it to a PyTorch tensor.
+        if isinstance(volume, np.ndarray):
+            return_numpy = True
+            volume = torch.from_numpy(volume)
+
+        # If running on CPU, ensure only a single volume is provided.
+        if self.device.type == 'cpu' and volume.dim() == 4:
+            raise AssertionError("Batched volumes are not allowed on CPU. Please provide a single volume.")
+
+        if volume.dim() == 4:
+            output = self.batched_rescale(volume)
+        else: 
+            output = self.single_rescale(volume)
+        
+        # Return to CPU if Compute is on GPU
+        if self.device == torch.device('cuda'):
+            output = output.cpu()
+            torch.cuda.empty_cache()
+
+        # Either return a numpy array or a torch tensor
+        if return_numpy:
+            return output.numpy()
+        else:
+            return output
+
+    def batched_rescale(self, volume: torch.Tensor):
+        """
+        Process a (batched) volume: move to device, perform FFT, crop in Fourier space,
+        and compute the inverse FFT.
+        """
+        volume = volume.to(self.device)
+        is_batched = (volume.dim() == 4)
+        if not is_batched:
+            volume = volume.unsqueeze(0)
+
+        fft_volume = torch.fft.fftn(volume, dim=(-3, -2, -1), norm='ortho')
+        fft_volume = torch.fft.fftshift(fft_volume, dim=(-3, -2, -1))
+        
+        start_d, start_h, start_w, new_depth, new_height, new_width = self.calculate_cropping(volume)
+        
+        fft_cropped = fft_volume[..., 
+                                 start_d:start_d + new_depth,
+                                 start_h:start_h + new_height,
+                                 start_w:start_w + new_width]
+        
+        fft_cropped = torch.fft.ifftshift(fft_cropped, dim=(-3, -2, -1))
+        out_volume = torch.fft.ifftn(fft_cropped, dim=(-3, -2, -1), norm='ortho')
+        out_volume = out_volume.real
+        
+        if not is_batched:
+            out_volume = out_volume.squeeze(0)
+            
+        return out_volume
+
+    def single_rescale(self, volume: torch.Tensor) -> torch.Tensor:
+        return self.batched_rescale(volume)
+
+    def calculate_cropping(self, volume: torch.Tensor):
+        """
+        Calculate cropping indices and new dimensions based on the voxel sizes.
+        """
+        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]
+        extent_width = in_width * self.input_voxel_size[2]
+        
+        new_depth = int(round(extent_depth / self.output_voxel_size[0]))
+        new_height = int(round(extent_height / self.output_voxel_size[1]))
+        new_width = int(round(extent_width / self.output_voxel_size[2]))
+        
+        # Ensure new dimensions are even
+        new_depth = new_depth - (new_depth % 2)
+        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)
+        
+        return start_d, start_h, start_w, new_depth, new_height, new_width

octopi/processing/evaluate.py

@@ -0,0 +1,289 @@
+from octopi import utils, io
+from scipy.spatial import distance
+from typing import List
+import copick, json, os
+import numpy as np
+
+class evaluator:
+
+    def __init__(self, 
+                 copick_config: str,
+                 ground_truth_user_id: str,
+                 ground_truth_session_id: str,
+                 prediction_user_id: str,
+                 predict_session_id: str,
+                 voxel_size: float = 10,
+                 beta: float = 4,
+                 object_names: List[str] = None):
+        
+        self.root = copick.from_file(copick_config)
+        print('Running Evaluation on the Following Copick Project: ', copick_config)
+
+        self.ground_truth_user_id = ground_truth_user_id
+        self.ground_truth_session_id = ground_truth_session_id
+        self.prediction_user_id = prediction_user_id
+        self.predict_session_id = predict_session_id
+        self.voxel_size = voxel_size
+        self.beta = beta
+        print(f'\nGround Truth Query: \nUserID: {ground_truth_user_id}, SessionID: {ground_truth_session_id}')
+        print(f'\nSubmitted Picks: \nUserID: {prediction_user_id}, SessionID: {predict_session_id}\n')
+
+        # Save input parameters
+        self.input_params = {
+            "copick_config": copick_config,
+            "ground_truth_user_id": ground_truth_user_id,
+            "ground_truth_session_id": ground_truth_session_id,
+            "prediction_user_id": prediction_user_id,
+            "predict_session_id": predict_session_id,
+        }
+
+        # Get objects that can be Picked
+        if not object_names:
+            print('No object names provided, using all pickable objects')
+            self.objects = [(obj.name, obj.radius) for obj in self.root.pickable_objects if obj.is_particle]
+        else:
+            # Get valid pickable objects with their radii
+            valid_objects = {obj.name: obj.radius for obj in self.root.pickable_objects if obj.is_particle}
+            
+            # Filter and validate provided object names
+            invalid_objects = [name for name in object_names if name not in valid_objects]
+            if invalid_objects:
+                print('WARNING: The following object names are not valid pickable objects:', invalid_objects)
+                print('Valid objects are:', list(valid_objects.keys()))
+            
+            self.objects = [(name, valid_objects[name]) for name in object_names if name in valid_objects]
+            
+            if not self.objects:
+                raise ValueError("None of the provided object names are valid pickable objects")
+                
+        print('Using the following valid objects:', [name for name, _ in self.objects])
+
+        # Define object-specific weights
+        self.weights = {
+            "apo-ferritin": 1,
+            "beta-amylase": 0,  # Excluded from scoring
+            "beta-galactosidase": 2,
+            "ribosome": 1,
+            "thyroglobulin": 2,
+            "virus-like particle": 1,
+        }
+
+    def run(self, 
+            save_path: str = None,
+            distance_threshold_scale: float = 0.8,
+            runIDs: List[str] = None):
+    
+        # Type check for runIDs
+        if runIDs is not None and not (isinstance(runIDs, list) and all(isinstance(x, str) for x in runIDs)):
+            raise TypeError("runIDs must be a list of strings")
+
+        run_ids = runIDs if runIDs else [run.name for run in self.root.runs]
+        print('\nRunning Metrics Evaluation on the Following RunIDs: ', run_ids)
+
+        metrics = {}
+        summary_metrics = {name: {'precision': [], 'recall': [], 'f1_score': [], 'fbeta_score': [], 'accuracy': [], 
+                                'true_positives': [], 'false_positives': [], 'false_negatives': []} for name, _ in self.objects}
+        
+        # For storing the aggregated counts per particle type (across all runs)
+        aggregated_counts = {name: {'total_tp': 0, 'total_fp': 0, 'total_fn': 0} for name, _ in self.objects}
+
+        for runID in run_ids:
+            # Initialize the nested dictionary for this runID
+            metrics[runID] = {}
+            run = self.root.get_run(runID)
+
+            for name, radius in self.objects:
+                
+                # Get Ground Truth and Predicted Coordinates
+                gt_coordinates = io.get_copick_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(
+                    run, name,
+                    self.prediction_user_id, self.predict_session_id, 
+                    self.voxel_size, raise_error=False
+                )
+                
+                # If no reference (GT) points, all candidate points are false positives
+                if gt_coordinates is None or len(gt_coordinates) == 0:
+                    num_pred_points = pred_coordinates.shape[0] if pred_coordinates is not None else 0
+                    metrics[runID][name] = {'precision': 0, 'recall': 0, 'fbeta_score': 0, 'true_positives': 0, 'false_positives': num_pred_points, 'false_negatives': 0}
+                    
+                    # Update aggregated counts
+                    aggregated_counts[name]['total_fp'] += num_pred_points
+                    
+                    continue
+
+                # If no candidate (predicted) points, all reference points are false negatives
+                if pred_coordinates is None or len(pred_coordinates) == 0:
+                    num_gt_points = gt_coordinates.shape[0] if gt_coordinates is not None else 0
+                    metrics[runID][name] = {'precision': 0, 'recall': 0, 'fbeta_score': 0, 'true_positives': 0, 'false_positives': 0, 'false_negatives': num_gt_points}
+                    
+                    # Update aggregated counts
+                    aggregated_counts[name]['total_fn'] += num_gt_points
+                    
+                    continue
+
+                # Compute Distance Threshold Based on Particle Radius
+                distance_threshold = (radius/self.voxel_size) * distance_threshold_scale
+                metrics[runID][name] = self.compute_metrics(gt_coordinates, pred_coordinates, distance_threshold)         
+
+                # Collect metrics for summary statistics
+                for key in summary_metrics[name]:
+                    summary_metrics[name][key].append(metrics[runID][name][key])
+                
+                # Update aggregated counts
+                aggregated_counts[name]['total_tp'] += metrics[runID][name]['true_positives']
+                aggregated_counts[name]['total_fp'] += metrics[runID][name]['false_positives']
+                aggregated_counts[name]['total_fn'] += metrics[runID][name]['false_negatives']
+
+        # Create a new dictionary for summarized metrics
+        final_summary_metrics = {}
+
+        # Compute average metrics and standard deviations across runs for each object
+        for name, _ in self.objects:
+            # Initialize the final summary for the object
+            final_summary_metrics[name] = {}
+
+            for key in summary_metrics[name]:
+                mu_val = float(np.mean(summary_metrics[name][key]))
+                std_val = float(np.std(summary_metrics[name][key]))
+
+                # Populate the new dictionary with structured data
+                final_summary_metrics[name][key] = {
+                    'mean': mu_val,
+                    'std': std_val
+                }
+
+        print('\nAverage Metrics Summary:')
+        self.print_metrics_summary(final_summary_metrics)          
+        
+        # Compute Final Kaggle Submission Score using reference approach
+        aggregate_fbeta = 0.0
+        total_weight = 0.0
+        
+        print('\nCalculating Final F-beta Score using per-particle approach:')
+        for name, counts in aggregated_counts.items():
+            tp = counts['total_tp']
+            fp = counts['total_fp']
+            fn = counts['total_fn']
+            
+            # Calculate precision and recall for this particle type
+            precision = tp / (tp + fp) if (tp + fp) > 0 else 0
+            recall = tp / (tp + fn) if (tp + fn) > 0 else 0
+            
+            # Calculate F-beta for this particle type
+            particle_fbeta = (1 + self.beta**2) * (precision * recall) / \
+                            ((self.beta**2 * precision) + recall) if \
+                            ((self.beta**2 * precision) + recall) > 0 else 0
+            
+            # Get the weight for this particle type
+            weight = self.weights.get(name, 1)
+            
+            # Accumulate weighted F-beta score
+            aggregate_fbeta += particle_fbeta * weight
+            total_weight += weight
+            
+            print(f"  {name}: TP={tp}, FP={fp}, FN={fn}, Precision={precision:.3f}, " + 
+                f"Recall={recall:.3f}, F-beta={particle_fbeta:.3f}, Weight={weight}")
+        
+        # Normalize by total weight
+        final_fbeta = aggregate_fbeta / total_weight if total_weight > 0 else 0
+        
+        print(f'\nFinal Kaggle Submission Score: {final_fbeta:.3f}')   
+
+        # Save average and detailed metrics with parameters included        
+        if save_path:
+            self.parameters = {
+                "distance_threshold_scale": distance_threshold_scale,
+                "runIDs": runIDs,
+            }    
+            
+            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")}')
+            
+            detailed_metrics = { "input": self.input_params, "parameters": self.parameters, 
+                                "metrics": metrics }
+            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")}')
+
+    def compute_metrics(self, 
+                        gt_points, 
+                        pred_points, 
+                        threshold):
+        
+        gt_points = np.array(gt_points)
+        pred_points = np.array(pred_points)
+        
+        # Calculate distances
+        if gt_points.shape[0] == 0:
+            # No ground truth points: all predictions are false positives
+            fp = pred_points.shape[0]
+            fn = 0
+            tp = 0
+        elif pred_points.shape[0] == 0:
+            # No predictions: all ground truth points are false negatives
+            fp = 0
+            fn = gt_points.shape[0]
+            tp = 0
+        else:    
+            # Calculate distances
+            dist_matrix = distance.cdist(pred_points, gt_points, 'euclidean')
+
+            # Determine matches within the threshold
+            tp = np.sum(np.min(dist_matrix, axis=1) < threshold)
+            fp = np.sum(np.min(dist_matrix, axis=1) >= threshold)
+            fn = np.sum(np.min(dist_matrix, axis=0) >= threshold)
+        
+        # Precision, Recall, F1 Score
+        precision = tp / (tp + fp) if tp + fp > 0 else 0
+        recall = tp / (tp + fn) if tp + fn > 0 else 0
+        f1_score = 2 * (precision * recall) / (precision + recall) if precision + recall > 0 else 0
+        accuracy = tp / (tp + fp + fn)  # Note: TN not considered here
+
+         # Compute F_beta using the formula
+        if (self.beta**2 * precision + recall) > 0:
+            fbeta = (1 + self.beta**2) * (precision * recall) / (self.beta**2 * precision + recall)
+        else:
+            fbeta = 0
+        
+        return {
+            'precision': precision,
+            'recall': recall,
+            'f1_score': f1_score,
+            'fbeta_score': fbeta,
+            'accuracy': accuracy, 
+            'true_positives': int(tp),
+            'false_positives': int(fp),
+            'false_negatives': int(fn)
+        }
+    
+    def print_metrics_summary(self, metrics_dict):
+        for name, metrics in metrics_dict.items():
+            recall = metrics['recall']
+            precision = metrics['precision']
+            f1_score = metrics['f1_score']
+            fbeta_score = metrics['fbeta_score']
+            false_positives = metrics['false_positives']
+            false_negatives = metrics['false_negatives']
+            
+            # Format the metrics for the current object
+            formatted_metrics = (
+                f"Recall: {recall['mean']:.3f} ± {recall['std']:.3f}, "
+                f"Precision: {precision['mean']:.3f} ± {precision['std']:.3f}, "
+                f"F1 Score: {f1_score['mean']:.3f} ± {f1_score['std']:.3f}, "
+                f"F_beta Score: {fbeta_score['mean']:.3f} ± {fbeta_score['std']:.3f}, "
+                f"False_Positives: {false_positives['mean']:.1f} ± {false_positives['std']:.1f}, "
+                f"False_Negatives: {false_negatives['mean']:.1f} ± {false_negatives['std']:.1f}"
+            )
+            
+            # Print the object name and its metrics
+            print(f"{name}: [{formatted_metrics}]")
+        print()

octopi/processing/importers.py

@@ -0,0 +1,213 @@
+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 tqdm import tqdm
+
+def from_dataportal(
+    config, 
+    datasetID,
+    overlay_path,
+    dataportal_name,
+    target_tomo_type,
+    input_voxel_size = 10,
+    output_voxel_size = None):
+    """
+    Download and process tomograms from the CZI Dataportal.
+    
+    Args:
+        config (str): Path to the copick configuration file
+        datasetID (int): ID of the dataset to download
+        overlay_path (str): Path to the overlay file
+        dataportal_name (str): Name of the tomogram type in the dataportal
+        target_tomo_alg (str): Name to use for the tomogram locally
+        input_voxel_size (float): Original voxel size of the tomograms
+        output_voxel_size (float, optional): Desired voxel size for downsampling
+    """
+    if config is not None:
+        root = copick.from_file(config)
+    elif datasetID is not None and overlay_path is not None:
+        root = copick.from_czcdp_datasets([datasetID], overlay_root=overlay_path)
+    else:
+        raise ValueError('Either config or datasetID and overlay_path must be provided')
+
+    # If we want to save the tomograms at a different voxel size, we need to rescale the tomograms
+    if output_voxel_size is not None and output_voxel_size > input_voxel_size:
+        rescale = FourierRescale(input_voxel_size, output_voxel_size)
+
+    # Create a directory for the tomograms
+    for run in tqdm(root.runs):
+
+        # Check if voxel spacing is available
+        vs = run.get_voxel_spacing(input_voxel_size)
+
+        if vs is None:
+            print(f'No Voxel-Spacing Available for RunID: {run.name}, Voxel-Size: {input_voxel_size}')
+            continue
+        
+        # Check if base reconstruction method is available
+        avail_tomos = vs.get_tomograms(dataportal_name)
+        if avail_tomos is None: 
+            print(f'No Tomograms Available for RunID: {run.name}, Voxel-Size: {input_voxel_size}, Tomo-Type: {dataportal_name}')
+            continue
+
+        # Download the tomogram
+        if len(avail_tomos) > 0:
+            vol = avail_tomos[0].numpy()
+
+            # 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)
+            else:
+                vol = rescale.run(vol)
+                write.tomogram(run, vol, output_voxel_size, target_tomo_type)
+    
+    print(f'Downloading Complete!! Downloaded {len(root.runs)} runs')
+
+def cli_dataportal_parser(parser_description, add_slurm: bool = False):
+    """
+    Create argument parser for the dataportal download command.
+    
+    Args:
+        parser_description (str): Description of the parser
+        add_slurm (bool): Whether to add SLURM-specific arguments
+    
+    Returns:
+        argparse.ArgumentParser: Configured argument parser
+    """
+    parser = argparse.ArgumentParser(
+        description=parser_description,
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    parser.add_argument('--config', type=str, required=False, default=None, help='Path to the config file')
+    parser.add_argument('--datasetID', type=int, required=False, default=None, help='Dataset ID')
+    parser.add_argument('--overlay-path', type=str, required=False, default=None, help='Path to the overlay file')
+    parser.add_argument('--dataportal-name', type=str, required=False, default='wbp', help='Dataportal name')
+    parser.add_argument('--target-tomo-type', type=str, required=False, default='wbp', help='Local name')
+    parser.add_argument('--input-voxel-size', type=float, required=False, default=10, help='Voxel size')
+    parser.add_argument('--output-voxel-size', type=float, required=False, default=None, help='Save voxel size')
+    
+    if add_slurm:
+        slurm_group = parser.add_argument_group("SLURM Arguments")
+        common.add_slurm_parameters(slurm_group, 'dataportal-importer', gpus = 0)
+
+    args = parser.parse_args()
+    return args
+
+def cli_dataportal():
+    """
+    Command-line interface for downloading tomograms from the Dataportal.
+    Handles argument parsing and calls from_dataportal with the parsed arguments.
+    """
+    parser_description = "Import tomograms from the Dataportal with optional downsampling with Fourier Cropping"
+    args = cli_dataportal_parser(parser_description)
+    from_dataportal(args.config, args.datasetID, args.overlay_path, args.dataportal_name, args.target_tomo_type, args.input_voxel_size, args.output_voxel_size)
+
+
+def from_mrcs(
+    mrcs_path,
+    config,
+    target_tomo_type,
+    input_voxel_size,
+    output_voxel_size = None):
+    """
+    Import and process tomograms from local MRC/MRCS files.
+    
+    Args:
+        mrcs_path (str): Path to directory containing MRC/MRCS files
+        config (str): Path to the copick configuration file
+        target_tomo_type (str): Name to use for the tomogram locally
+        input_voxel_size (float): Original voxel size of the tomograms
+        output_voxel_size (float, optional): Desired voxel size for downsampling
+    """
+    # Load Copick Project
+    if os.path.exists(config):
+        root = copick.from_file(config)
+    else:
+        raise ValueError('Config file not found')
+
+    # List all .mrc and .mrcs files in the directory
+    mrc_files = glob.glob(os.path.join(mrcs_path, "*.mrc")) + glob.glob(os.path.join(mrcs_path, "*.mrcs"))
+    if not mrc_files:
+        print(f"No .mrc or .mrcs files found in {mrcs_path}")
+        return
+
+    # Prepare rescaler if needed
+    rescale = None
+    if output_voxel_size is not None and output_voxel_size > input_voxel_size:
+        rescale = FourierRescale(input_voxel_size, output_voxel_size)        
+
+    # Check if the mrcs file exists
+    if not os.path.exists(mrcs_path):
+        raise FileNotFoundError(f'MRCs file not found: {mrcs_path}')
+    
+    for mrc_path in tqdm(mrc_files):
+
+        # Get or Create Run
+        runID = os.path.splitext(os.path.basename(mrc_path))[0]
+        try:
+            run = root.new_run(runID)
+        except Exception as e:
+            run = root.get_run(runID)
+
+        # Load the mrcs file
+        with mrcfile.open(mrc_path) as mrc:
+            vol = mrc.data
+            # Check voxel size in MRC header vs user input
+            mrc_voxel_size = float(mrc.voxel_size.x)  # assuming cubic voxels
+            if abs(mrc_voxel_size - input_voxel_size) > 1e-1:
+                print(f"WARNING: Voxel size in {mrc_path} header ({mrc_voxel_size}) "
+                      f"differs from user input ({input_voxel_size})")
+
+        # Rescale if needed
+        if rescale is not None:
+            vol = rescale.run(vol)
+            voxel_size_to_write = output_voxel_size
+        else:
+            voxel_size_to_write = input_voxel_size
+
+        # Write the tomogram
+        write.tomogram(run, vol, voxel_size_to_write, target_tomo_type)
+    print(f"Processed {len(mrc_files)} files from {mrcs_path}")
+
+
+def cli_mrcs_parser(parser_description, add_slurm: bool = False):
+    """
+    Create argument parser for the MRC import command.
+    
+    Args:
+        parser_description (str): Description of the parser
+        add_slurm (bool): Whether to add SLURM-specific arguments
+    
+    Returns:
+        argparse.ArgumentParser: Configured argument parser
+    """
+    parser = argparse.ArgumentParser(
+        description=parser_description,
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    # Input Arguments
+    parser.add_argument('--mrcs-path', type=str, required=True, help='Path to the mrcs file')
+    parser.add_argument('--config', type=str, required=False, default=None, help='Path to the config file to write tomograms to')
+    parser.add_argument('--target-tomo-type', type=str, required=True, help='Reconstruction algorithm used to create the tomogram')
+    parser.add_argument('--input-voxel-size', type=float, required=False, default=10, help='Voxel size of the MRC tomogram')
+    parser.add_argument('--output-voxel-size', type=float, required=False, default=None, help='Output voxel size (if desired to downsample to lower resolution)')
+    
+    if add_slurm:
+        slurm_group = parser.add_argument_group("SLURM Arguments")
+        common.add_slurm_parameters(slurm_group, 'mrcs-importer', gpus = 0)
+
+    args = parser.parse_args()
+
+    return args
+    
+def cli_mrcs():
+    """
+    Command-line interface for importing MRC/MRCS files.
+    Handles argument parsing and calls from_mrcs with the parsed arguments.
+    """
+    parser_description = "Import MRC volumes from a directory."
+    args = cli_mrcs_parser(parser_description)
+    from_mrcs(args.mrcs_path, args.config, args.target_tomo_type, args.input_voxel_size, args.output_voxel_size)

octopi/processing/my_metrics.py

@@ -0,0 +1,26 @@
+from monai.utils import (MetricReduction, look_up_option)
+from monai.metrics import confusion_matrix as monai_cm
+from typing import Any
+import torch, mlflow
+
+def my_log_param(params_dict, client = None, trial_run_id = None):
+
+    if client is not None and trial_run_id is not None:
+        # client.log_params(run_id=trial_run_id, params=params_dict)
+        for key, value in params_dict.items():
+            client.log_param(run_id=trial_run_id, key=key, value=value)
+    else:
+        mlflow.log_params(params_dict)
+
+
+##############################################################################################################################
+
+def my_log_metric(metric_name, val, curr_step, client = None, trial_run_id = None):
+
+    if client is not None and trial_run_id is not None:
+        client.log_metric(run_id = trial_run_id, 
+                          key = metric_name,
+                          value = val, 
+                          step = curr_step)
+    else:
+        mlflow.log_metric(metric_name, val, step = curr_step)