octopi 1.4.0__py3-none-any.whl

octopi/processing/create_targets_from_picks.py

@@ -0,0 +1,224 @@
+from octopi.processing.segmentation_from_picks import from_picks
+from copick_utils.io import readers, writers
+import zarr, os, yaml, copick
+from octopi.utils import io
+from typing import List
+from tqdm import tqdm
+import numpy as np
+
+def print_target_summary(train_targets: dict, target_segmentation_name: str, maxval: int):
+    """
+    Print a summary of the target volume structure.
+    """
+    print("\n" + "="*60)
+    print("TARGET VOLUME SUMMARY")
+    print("="*60)
+    print(f"Segmentation name: {target_segmentation_name}")
+    print(f"Total classes: {len(train_targets) + 1} (including background)")
+    print("\nLabel Index → Object Name (Type):")
+    print(f"  {0:3d} → background")
+    
+    # Sort by label for display
+    sorted_targets = sorted(train_targets.items(), key=lambda x: x[1]['label'])
+    for name, info in sorted_targets:
+        obj_type = "particle" if info['is_particle_target'] else "segmentation"
+        radius_info = f", radius={info['radius']:.1f}Å" if info['radius'] else ""
+        print(f"  {info['label']:3d} → {name} ({obj_type}{radius_info})")
+    
+    print("="*60)
+    print(f"💡 Use --num-classes {maxval + 1} when training with this target")
+    print("="*60 + "\n")
+
+def generate_targets(
+    config,
+    train_targets: dict,
+    voxel_size: float = 10,
+    tomo_algorithm: str = 'wbp',
+    radius_scale: float = 0.8,    
+    target_segmentation_name: str = 'targets',
+    target_user_name: str = 'octopi',
+    target_session_id: str = '1',
+    run_ids: List[str] = None,
+    ):
+    """
+    Generate segmentation targets from picks in CoPick configuration.
+
+    Args:
+        copick_config_path (str): Path to CoPick configuration file.
+        picks_user_id (str): User ID associated with picks.
+        picks_session_id (str): Session ID associated with picks.
+        target_segmentation_name (str): Name for the target segmentation.
+        target_user_name (str): User name associated with target segmentation.
+        target_session_id (str): Session ID for the target segmentation.
+        voxel_size (float): Voxel size for tomogram reconstruction.
+        tomo_algorithm (str): Tomogram reconstruction algorithm.
+        radius_scale (float): Scale factor for target object radius.
+    """
+
+    # Default session ID to 1 if not provided
+    root = copick.from_file(config)
+    if target_session_id is None:
+        target_session_id = '1'
+
+    # Print target summary
+    print('🔄 Creating Targets for the following objects:', ', '.join(train_targets.keys()))
+
+    # Get Target Names
+    target_names = list(train_targets.keys())
+
+    # If runIDs are not provided, load all runs
+    if run_ids is None:
+        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
+    maxval = -1
+    for runID in tqdm(run_ids):
+
+        # Get Run
+        numPicks = 0
+        run = root.get_run(runID)
+
+        # Get Target Shape
+        vs = run.get_voxel_spacing(voxel_size)
+        if vs is None:
+            print(f"⚠️ Warning: skipping run {runID} with no voxel spacing {voxel_size}")
+            continue
+        tomo = vs.get_tomogram(tomo_algorithm)
+        if tomo is None:
+            print(f"⚠️ Warning: skipping run {runID} with no tomogram {tomo_algorithm}")
+            continue
+        
+        # Initialize Target Volume
+        loc = tomo.zarr()
+        shape = zarr.open(loc)['0'].shape
+        target = np.zeros(shape, dtype=np.uint8)
+
+        # Generate Targets
+        # Applicable segmentations
+        query_seg = []
+        for target_name in target_names:
+            if not train_targets[target_name]["is_particle_target"]:            
+                query_seg += run.get_segmentations(
+                    name=target_name,
+                    user_id=train_targets[target_name]["user_id"],
+                    session_id=train_targets[target_name]["session_id"],
+                    voxel_size=voxel_size
+                )
+
+        # Add Segmentations to Target
+        for seg in query_seg:
+            classLabel = train_targets[seg.name]['label']
+            segvol = seg.numpy()
+            # Set all non-zero values to the class label
+            segvol[segvol > 0] = classLabel
+            target = np.maximum(target, segvol)
+
+        # Applicable picks
+        query = []
+        for target_name in target_names:
+            if train_targets[target_name]["is_particle_target"]:
+                query += run.get_picks(
+                    object_name=target_name,
+                    user_id=train_targets[target_name]["user_id"],
+                    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)
+            target = from_picks(pick, 
+                                target, 
+                                train_targets[pick.pickable_object_name]['radius'] * radius_scale,
+                                train_targets[pick.pickable_object_name]['label'],
+                                voxel_size
+                                )
+
+        # Write Segmentation for non-empty targets
+        if target.max() > 0:
+            tqdm.write(f'📝 Annotating {numPicks} picks in {runID}...')    
+            writers.segmentation(run, target, target_user_name, 
+                               name = target_segmentation_name, session_id= target_session_id, 
+                               voxel_size = voxel_size)
+        if target.max() > maxval:
+            maxval = target.max()
+    
+    print('✅ Creation of targets complete!')
+
+    # Save Parameters
+    overlay_root = io.remove_prefix(root.config.overlay_root)
+    basepath = os.path.join(overlay_root, 'logs')
+    os.makedirs(basepath, exist_ok=True)
+    labels = {name: info['label'] for name, info in train_targets.items()}
+    args = {
+        "config": config,
+        "train_targets": train_targets,
+        "radius_scale": radius_scale,
+        "tomo_algorithm": tomo_algorithm,
+        "target_name": target_segmentation_name,
+        "target_user_name": target_user_name,
+        "target_session_id": target_session_id,
+        "voxel_size": voxel_size,
+        "labels": labels,
+    }
+    target_query = f'{target_user_name}_{target_session_id}_{target_segmentation_name}'
+    print(f'💾 Saving parameters to {basepath}/targets-{target_query}.yaml')
+    save_parameters(args, basepath, target_query)
+
+    # Print Target Summary
+    print_target_summary(train_targets, target_segmentation_name, maxval)
+
+def save_parameters(args, basepath: str, target_query: str):
+    """
+    Save parameters to a YAML file with subgroups for input, output, and parameters.
+    Append to the file if it already exists.
+
+    Args:
+        args: Parsed arguments from argparse.
+        basepath: Path to save the YAML file.
+        target_query: Query string for target identification.
+    """
+    # Prepare input group
+    keys = ['user_id', 'session_id']
+    input_group = {
+        "config": args['config'],
+        "labels": {name: info['label'] for name, info in args['train_targets'].items()},  # <-- Added comma here
+        "targets": {name: {k: info[k] for k in keys} for name, info in args['train_targets'].items()}
+    }
+        
+    # Organize parameters into subgroups
+    new_entry = {
+        "input": input_group,
+        "parameters": {
+            "radius_scale": args["radius_scale"],
+            "tomogram_algorithm": args["tomo_algorithm"],
+            "voxel_size": args["voxel_size"],
+        }
+    }
+
+    # Check if the YAML file already exists
+    output_path = os.path.join(
+        basepath, 
+        f'targets-{args["target_user_name"]}_{args["target_session_id"]}_{args["target_name"]}.yaml')
+    if os.path.exists(output_path):
+        # Load the existing content
+        with open(output_path, 'r') as f:
+            try:
+                existing_data = yaml.safe_load(f)
+                if existing_data is None:
+                    existing_data = {}  # Ensure it's a dictionary
+                elif not isinstance(existing_data, dict):
+                    raise ValueError("Existing YAML data is not a dictionary. Cannot update.")
+            except yaml.YAMLError:
+                existing_data = {}  # Treat as empty if the file is malformed
+    else:
+        existing_data = {}  # Initialize as empty dictionary if the file does not exist
+
+    # Save back to the YAML file
+    io.save_parameters_yaml(new_entry, output_path)

octopi/processing/downloader.py

@@ -0,0 +1,138 @@
+import rich_click as click
+
+def from_dataportal(
+    config, 
+    datasetID,
+    overlay_path,
+    source_type,
+    target_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
+        source_type (str): Name of the tomogram type in the dataportal
+        target_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
+    """
+
+    from octopi.processing.downsample import FourierRescale
+    from octopi.utils.progress import _progress, print_summary
+    from copick_utils.io import writers
+    import copick   
+
+    # Either load an existing configuration file or create one from datasetID and overlay_path
+    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, 
+            output_path='config.json', overlay_fs_args={'auto_mkdir': True}
+        )
+    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)
+    else:
+        output_voxel_size = None
+
+    # Print Parameter Summary
+    print_summary(
+        "Download Tomograms",
+        datasetID=datasetID, overlay_path=overlay_path,
+        config=config, source_type=source_type, target_type=target_type,
+        input_voxel_size=input_voxel_size, output_voxel_size=output_voxel_size,
+    )
+
+    # Main Loop
+    for run in _progress(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(source_type)
+        if avail_tomos is None: 
+            print(f'No Tomograms Available for RunID: {run.name}, Voxel-Size: {input_voxel_size}, Tomo-Type: {source_type}')
+            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:
+                writers.tomogram(run, vol, input_voxel_size, target_type)
+            else:
+                vol = rescale.run(vol)
+                writers.tomogram(run, vol, output_voxel_size, target_type)
+    
+    print(f'✅ Download Complete!\nDownloaded {len(root.runs)} runs')
+
+
+@click.command('download')
+# Voxel Settings
+@click.option('-ovs', '--output-voxel-size', type=float, default=None,
+              help="Desired output voxel size for downsampling (optional)")
+@click.option('-ivs', '--input-voxel-size', type=float, default=10,
+              help="Original voxel size of the tomograms")
+# Tomogram Settings
+@click.option('-t', '--target-type', type=str, default='denoised',
+              help="Local tomogram type name to save in your Copick project.")
+@click.option('-s', '--source-type', type=str, default='wbp-denoised-ctfdeconv',
+              help="Name of the tomogram type as labeled on the CryoET Data Portal")
+# Input Arguments
+@click.option('-o', '--overlay', type=click.Path(), default=None,
+              help="Path to the overlay directory (required with datasetID)")
+@click.option('-ds', '--datasetID', type=int, default=None,
+              help="Dataset ID from CZI Dataportal (alternative to config)")
+@click.option('-c', '--config', type=click.Path(exists=True), default=None,
+              help="Path to the copick configuration file (alternative to datasetID)")
+def cli(config, datasetid, overlay, source_type, target_type,
+        input_voxel_size, output_voxel_size):
+    """
+    Download and (optionally) downsample tomograms from the CryoET-DataPortal.
+    
+    This command fetches reconstructed tomograms from publicly available datasets and saves them 
+    to your local copick project. Downsampling is performed via Fourier cropping to preserve 
+    high-frequency information while reducing file size.
+    
+    Two modes of operation:
+    
+    \b
+    1. Download tomograms and downsample to a new voxel size:
+       octopi download -c config.json --input-voxel-size 10 --output-voxel-size 20
+    
+    \b
+    2. Create new project from the CryoET Data Portal:
+       octopi download --datasetID 10301 --overlay-path ./my_project --input-voxel-size 10
+    
+    The downloaded tomograms will be stored in your copick project structure with the specified 
+    voxel spacing and tomogram type.
+    """
+    
+    from_dataportal(
+        config=config,
+        datasetID=datasetid,
+        overlay_path=overlay,
+        source_type=source_type,
+        target_type=target_type,
+        input_voxel_size=input_voxel_size,
+        output_voxel_size=output_voxel_size
+    )
+
+
+if __name__ == "__main__":
+    cli()
+

octopi/processing/downsample.py

@@ -0,0 +1,125 @@
+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:]
+        
+        # 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 - 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