segger.prediction

prediction module for Segger.

Contains the implementation of the Segger model using Graph Neural Networks.

Prediction module for Segger.

Contains prediction scripts and utilities for the Segger model.

load_model

load_model(checkpoint_path)

Load a LitSegger model from a checkpoint.

Parameters

checkpoint_path : str Specific checkpoint file to load, or directory where the model checkpoints are stored. If directory, the latest checkpoint is loaded.

Returns

LitSegger The loaded LitSegger model.

Raises

FileNotFoundError If the specified checkpoint file does not exist.

Source code in src/segger/prediction/predict_parquet.py

def load_model(checkpoint_path: str) -> LitSegger:
    """
    Load a LitSegger model from a checkpoint.

    Parameters
    ----------
    checkpoint_path : str
        Specific checkpoint file to load, or directory where the model checkpoints are stored.
        If directory, the latest checkpoint is loaded.

    Returns
    -------
    LitSegger
        The loaded LitSegger model.

    Raises
    ------
    FileNotFoundError
        If the specified checkpoint file does not exist.
    """
    checkpoint_path = Path(checkpoint_path)
    msg = f"No checkpoint found at {checkpoint_path}. Please make sure you've provided the correct path."

    # Get last checkpoint if directory is provided
    if os.path.isdir(checkpoint_path):
        checkpoints = glob.glob(str(checkpoint_path / "*.ckpt"))
        if len(checkpoints) == 0:
            raise FileNotFoundError(msg)

        # Sort checkpoints by epoch and step
        def sort_order(c):
            match = re.match(r".*epoch=(\d+)-step=(\d+).ckpt", c)
            return int(match[1]), int(match[2])

        checkpoint_path = Path(sorted(checkpoints, key=sort_order)[-1])
    elif not checkpoint_path.exists():
        raise FileExistsError(msg)

    # Load model from checkpoint
    lit_segger = LitSegger.load_from_checkpoint(
        checkpoint_path=checkpoint_path,
    )

    return lit_segger

segment

segment(model, dm, save_dir, seg_tag, transcript_file, score_cut=0.5, use_cc=True, file_format='', save_transcripts=True, save_anndata=True, save_cell_masks=False, receptive_field={'k_bd': 4, 'dist_bd': 10, 'k_tx': 5, 'dist_tx': 3}, knn_method='cuda', verbose=False, gpu_ids=['0'], **anndata_kwargs)

Perform segmentation using the model, save transcripts, AnnData, and cell masks as needed, and log the parameters used during segmentation.

Parameters:

Name	Type	Description	Default
model	LitSegger	The trained segmentation model.	required
dm	SeggerDataModule	The SeggerDataModule instance for data loading.	required
save_dir	Union[str, Path]	Directory to save the final segmentation results.	required
seg_tag	str	Tag to include in the saved filename.	required
transcript_file	Union[str, Path]	Path to the transcripts Parquet file.	required
score_cut	float	The threshold for assigning transcripts to cells based on similarity scores. Defaults to 0.5.	0.5
use_cc	bool	If True, perform connected components analysis for unassigned transcripts. Defaults to True.	True
save_transcripts	bool	Whether to save the transcripts as Parquet. Defaults to True.	True
save_anndata	bool	Whether to save the results in AnnData format. Defaults to True.	True
save_cell_masks	bool	Save cell masks as Dask Geopandas Parquet. Defaults to False.	False
receptive_field	dict	Defines the receptive field for transcript-cell and transcript-transcript relations. Defaults to {'k_bd': 4, 'dist_bd': 10, 'k_tx': 5, 'dist_tx': 3}.	{'k_bd': 4, 'dist_bd': 10, 'k_tx': 5, 'dist_tx': 3}
knn_method	str	The method to use for nearest neighbors ('cuda' or 'kd_tree'). Defaults to 'cuda'.	'cuda'
verbose	bool	Whether to print verbose status updates. Defaults to False.	False
**anndata_kwargs		Additional keyword arguments passed to the create_anndata function.	{}

Returns:

Type	Description
None	None. Saves the result to disk in various formats and logs the parameter choices.

Source code in src/segger/prediction/predict_parquet.py

def segment(
    model: LitSegger,
    dm: SeggerDataModule,
    save_dir: Union[str, Path],
    seg_tag: str,
    transcript_file: Union[str, Path],
    score_cut: float = 0.5,
    use_cc: bool = True,
    file_format: str = "",
    save_transcripts: bool = True,
    save_anndata: bool = True,
    save_cell_masks: bool = False,  # Placeholder for future implementation
    receptive_field: dict = {"k_bd": 4, "dist_bd": 10, "k_tx": 5, "dist_tx": 3},
    knn_method: str = "cuda",
    verbose: bool = False,
    gpu_ids: list = ["0"],
    **anndata_kwargs,
) -> None:
    """
    Perform segmentation using the model, save transcripts, AnnData, and cell masks as needed,
    and log the parameters used during segmentation.

    Args:
        model (LitSegger): The trained segmentation model.
        dm (SeggerDataModule): The SeggerDataModule instance for data loading.
        save_dir (Union[str, Path]): Directory to save the final segmentation results.
        seg_tag (str): Tag to include in the saved filename.
        transcript_file (Union[str, Path]): Path to the transcripts Parquet file.
        score_cut (float, optional): The threshold for assigning transcripts to cells based on
                                     similarity scores. Defaults to 0.5.
        use_cc (bool, optional): If True, perform connected components analysis for unassigned
                                 transcripts. Defaults to True.
        save_transcripts (bool, optional): Whether to save the transcripts as Parquet. Defaults to True.
        save_anndata (bool, optional): Whether to save the results in AnnData format. Defaults to True.
        save_cell_masks (bool, optional): Save cell masks as Dask Geopandas Parquet. Defaults to False.
        receptive_field (dict, optional): Defines the receptive field for transcript-cell and
                                          transcript-transcript relations. Defaults to
                                          {'k_bd': 4, 'dist_bd': 10, 'k_tx': 5, 'dist_tx': 3}.
        knn_method (str, optional): The method to use for nearest neighbors ('cuda' or 'kd_tree').
                                    Defaults to 'cuda'.
        verbose (bool, optional): Whether to print verbose status updates. Defaults to False.
        **anndata_kwargs: Additional keyword arguments passed to the `create_anndata` function.

    Returns:
        None. Saves the result to disk in various formats and logs the parameter choices.
    """

    start_time = time()

    # Create a subdirectory with important parameter info (receptive field values)
    sub_dir_name = f"{seg_tag}_{score_cut}_{use_cc}_{receptive_field['k_bd']}_{receptive_field['dist_bd']}_{receptive_field['k_tx']}_{receptive_field['dist_tx']}_{datetime.now().strftime('%Y%m%d')}"
    save_dir = Path(save_dir) / sub_dir_name
    save_dir.mkdir(parents=True, exist_ok=True)

    # Paths for saving the output_ddf and edge_index Parquet files
    output_ddf_save_path = save_dir / "transcripts_df.parquet"
    edge_index_save_path = save_dir / "edge_index.parquet"

    if output_ddf_save_path.exists():
        warnings.warn(f"Removing existing file: {output_ddf_save_path}")
        shutil.rmtree(output_ddf_save_path)

    if use_cc:
        if edge_index_save_path.exists():
            warnings.warn(f"Removing existing file: {edge_index_save_path}")
            shutil.rmtree(edge_index_save_path)

    if verbose:
        print(f"Starting segmentation for {seg_tag}...")

    # Step 1: Load the data loaders from the SeggerDataModule
    step_start_time = time()
    train_dataloader = dm.train_dataloader()
    val_dataloader = dm.val_dataloader()
    test_dataloader = dm.test_dataloader()

    # Loop through the data loaders (train, val, and test)
    for loader_name, loader in zip(
        ["Train", "Validation", "Test"], [train_dataloader, val_dataloader, test_dataloader]
    ):
        # for loader_name, loader in zip(['Test'], [test_dataloader]):
        if verbose:
            print(f"Processing {loader_name} data...")

        for batch in tqdm(loader, desc=f"Processing {loader_name} batches"):
            gpu_id = random.choice(gpu_ids)
            # Call predict_batch for each batch
            predict_batch(
                model,
                batch,
                score_cut,
                receptive_field,
                use_cc=use_cc,
                knn_method=knn_method,
                edge_index_save_path=edge_index_save_path,
                output_ddf_save_path=output_ddf_save_path,
                gpu_id=gpu_id,
            )

    if verbose:
        elapsed_time = time() - step_start_time
        print(f"Batch processing completed in {elapsed_time:.2f} seconds.")

    seg_final_dd = pd.read_parquet(output_ddf_save_path)
    seg_final_dd = seg_final_dd.set_index("transcript_id")

    step_start_time = time()
    if verbose:
        print(f"Applying max score selection logic...")

    # Step 1: Find max bound indices (bound == 1) and max unbound indices (bound == 0)
    max_bound_idx = seg_final_dd[seg_final_dd["bound"] == 1].groupby("transcript_id")["score"].idxmax()
    max_unbound_idx = seg_final_dd[seg_final_dd["bound"] == 0].groupby("transcript_id")["score"].idxmax()

    # Step 2: Combine indices, prioritizing bound=1 scores
    final_idx = max_bound_idx.combine_first(max_unbound_idx)

    # Step 3: Use the computed final_idx to select the best assignments
    # Make sure you are using the divisions and set the index correctly before loc
    seg_final_filtered = seg_final_dd.loc[final_idx]

    if verbose:
        elapsed_time = time() - step_start_time
        print(f"Max score selection completed in {elapsed_time:.2f} seconds.")

    # Step 3: Load the transcripts DataFrame and merge results

    if verbose:
        print(f"Loading transcripts from {transcript_file}...")

    transcripts_df = pd.read_parquet(transcript_file)
    transcripts_df["transcript_id"] = transcripts_df["transcript_id"].astype(str)

    step_start_time = time()
    if verbose:
        print(f"Merging segmentation results with transcripts...")

    # Outer merge to include all transcripts, even those without assigned cell ids
    transcripts_df_filtered = transcripts_df.merge(seg_final_filtered, on="transcript_id", how="outer")
    if verbose:
        elapsed_time = time() - step_start_time
        print(f"Merged segmentation results with transcripts in {elapsed_time:.2f} seconds.")

    if use_cc:

        step_start_time = time()
        if verbose:
            print(f"Computing connected components for unassigned transcripts...")
        # Load edge indices from saved Parquet
        edge_index_dd = pd.read_parquet(edge_index_save_path)

        # Step 2: Get unique transcript_ids from edge_index_dd and their positional indices
        transcript_ids_in_edges = pd.concat([edge_index_dd["source"], edge_index_dd["target"]]).unique()

        # Create a lookup table with unique indices
        lookup_table = pd.Series(data=range(len(transcript_ids_in_edges)), index=transcript_ids_in_edges).to_dict()

        # Map source and target to positional indices
        edge_index_dd["index_source"] = edge_index_dd["source"].map(lookup_table)
        edge_index_dd["index_target"] = edge_index_dd["target"].map(lookup_table)
        # Step 3: Compute connected components for transcripts involved in edges
        source_indices = np.asarray(edge_index_dd["index_source"])
        target_indices = np.asarray(edge_index_dd["index_target"])
        data_cp = np.ones(len(source_indices), dtype=np.float32)

        # Create the sparse COO matrix
        coo_cp_matrix = scipy_coo_matrix(
            (data_cp, (source_indices, target_indices)),
            shape=(len(transcript_ids_in_edges), len(transcript_ids_in_edges)),
        )

        # Use CuPy's connected components algorithm to compute components
        n, comps = cc(coo_cp_matrix, directed=True, connection="strong")
        if verbose:
            elapsed_time = time() - step_start_time
            print(f"Computed connected components for unassigned transcripts in {elapsed_time:.2f} seconds.")

        step_start_time = time()
        if verbose:
            print(f"The rest...")
        # # Step 4: Map back the component labels to the original transcript_ids

        def _get_id():
            """Generate a random Xenium-style ID."""
            return "".join(np.random.choice(list("abcdefghijklmnopqrstuvwxyz"), 8)) + "-nx"

        new_ids = np.array([_get_id() for _ in range(n)])
        comp_labels = new_ids[comps]
        comp_labels = pd.Series(comp_labels, index=transcript_ids_in_edges)
        # Step 5: Handle only unassigned transcripts in transcripts_df_filtered
        unassigned_mask = transcripts_df_filtered["segger_cell_id"].isna()

        unassigned_transcripts_df = transcripts_df_filtered.loc[unassigned_mask, ["transcript_id"]]

        # Step 6: Map component labels only to unassigned transcript_ids
        new_segger_cell_ids = unassigned_transcripts_df["transcript_id"].map(comp_labels)

        # Step 7: Create a DataFrame with updated 'segger_cell_id' for unassigned transcripts
        unassigned_transcripts_df = unassigned_transcripts_df.assign(segger_cell_id=new_segger_cell_ids)

        # Step 8: Merge this DataFrame back into the original to update only the unassigned segger_cell_id

        # Merging the updates back to the original DataFrame
        transcripts_df_filtered = transcripts_df_filtered.merge(
            unassigned_transcripts_df[["transcript_id", "segger_cell_id"]],
            on="transcript_id",
            how="left",  # Perform a left join to only update the unassigned rows
            suffixes=("", "_new"),  # Suffix for new column to avoid overwriting
        )

        # Step 9: Fill missing segger_cell_id values with the updated values from the merge
        transcripts_df_filtered["segger_cell_id"] = transcripts_df_filtered["segger_cell_id"].fillna(
            transcripts_df_filtered["segger_cell_id_new"]
        )

        transcripts_df_filtered = transcripts_df_filtered.drop(columns=["segger_cell_id_new"])

        if verbose:
            elapsed_time = time() - step_start_time
            print(f"The rest computed in {elapsed_time:.2f} seconds.")

    # Step 5: Save the merged results based on options

    if save_transcripts:
        if verbose:
            step_start_time = time()
            print(f"Saving transcripts.parquet...")
        transcripts_save_path = save_dir / "segger_transcripts.parquet"
        # transcripts_df_filtered = transcripts_df_filtered.repartition(npartitions=100)
        transcripts_df_filtered.to_parquet(
            transcripts_save_path,
            engine="pyarrow",  # PyArrow is faster and recommended
            compression="snappy",  # Use snappy compression for speed
            # write_index=False,  # Skip writing index if not needed
            # append=False,  # Set to True if you're appending to an existing Parquet file
            # overwrite=True,
        )  # Dask handles Parquet well
        if verbose:
            elapsed_time = time() - step_start_time
            print(f"Saved trasncripts.parquet in {elapsed_time:.2f} seconds.")

    if save_anndata:
        if verbose:
            step_start_time = time()
            print(f"Saving anndata object...")
        anndata_save_path = save_dir / "segger_adata.h5ad"
        segger_adata = create_anndata(transcripts_df_filtered, **anndata_kwargs)  # Compute for AnnData
        segger_adata.write(anndata_save_path)
        if verbose:
            elapsed_time = time() - step_start_time
            print(f"Saved anndata object in {elapsed_time:.2f} seconds.")

    if save_cell_masks:
        if verbose:
            step_start_time = time()
            print(f"Computing and saving cell masks anndata object...")
        # Placeholder for future cell masks implementation as Dask Geopandas Parquet
        boundaries_gdf = generate_boundaries(transcripts_df_filtered)
        cell_masks_save_path = save_dir / "segger_cell_boundaries.parquet"

        boundaries_gdf.to_parquet(cell_masks_save_path)
        if verbose:
            elapsed_time = time() - step_start_time
            print(f"Saved cell masks in {elapsed_time:.2f} seconds.")

    if verbose:
        elapsed_time = time() - step_start_time
        print(f"Results saved in {elapsed_time:.2f} seconds at {save_dir}.")

    # Step 6: Save segmentation parameters as a JSON log
    log_data = {
        "seg_tag": seg_tag,
        "score_cut": score_cut,
        "use_cc": use_cc,
        "receptive_field": receptive_field,
        "knn_method": knn_method,
        "save_transcripts": save_transcripts,
        "save_anndata": save_anndata,
        "save_cell_masks": save_cell_masks,
        "timestamp": datetime.now().isoformat(),
    }

    log_path = save_dir / "segmentation_log.json"
    with open(log_path, "w") as log_file:
        json.dump(log_data, log_file, indent=4)

    # Step 7: Garbage collection and memory cleanup
    torch.cuda.empty_cache()
    gc.collect()

    # Total time taken for the segmentation process
    if verbose:
        total_time = time() - start_time
        print(f"Total segmentation process completed in {total_time:.2f} seconds.")