update source code

src/build/lib/loki/decompose.py

@@ -77,11 +77,11 @@ def cell_type_decompose(sc_ad, st_ad, cell_type_col='cell_type', NMS_mode=False,
     
     :param sc_ad: AnnData object containing single-cell meta data.
     :param st_ad: AnnData object containing spatial data (ST or image) meta data.
-    :param density_prior: A numpy array providing prior information about cell densities in spatial spots.
     :param cell_type_col: The column name in `sc_ad.obs` that contains cell type annotations. Default is 'cell_type'.
-    :param target_count: If True, sums up the total number of cells in `st_ad.obs['cell_num']`. Can also be set to a specific value.
-    :param pca_mode: Boolean flag to apply PCA for dimensionality reduction. Default is True.
-    :param n_components: Number of PCA components to use if `pca_mode` is True. Default is 300.
+    :param NMS_mode: Boolean flag to apply Non-Maximum Suppression (NMS) mode. Default is False.
+    :param major_types: Major cell types used for NMS mode. Default is None.
+    :param min_percentile: The lower percentile used for clipping (defaults to 5).
+    :param max_percentile: The upper percentile used for clipping (defaults to 95).
     :return: The spatial AnnData object with projected cell type annotations.
     """
     

src/build/lib/loki/plot.py

@@ -8,107 +8,102 @@ import numpy as np
 from tqdm import tqdm
 
 
-
-def plot_alignment(ad_tar_coor, ad_src_coor, homo_coor, pca_hex_comb, tar_features, shift=300, s=0.8, boundary_line=True):
+def plot_alignment(
+    ad_tar_coor: np.ndarray,
+    ad_src_coor: np.ndarray,
+    homo_coor: np.ndarray,
+    pca_hex_comb: np.ndarray,
+    tar_features: np.ndarray,
+    shift: float = 300,
+    s: float = 0.8,
+    boundary_line: bool = True
+) -> None:
     """
-    Plots the target coordinates and alignment of source coordinates.
-
-    :param ad_tar_coor: Numpy array of target coordinates to be plotted in the first subplot.
-    :param ad_src_coor: Numpy array of source coordinates to be plotted in the second subplot.
-    :param homo_coor: Numpy array of alignment of source coordinates to be plotted in the third subplot.
-    :param pca_hex_comb: Color values (e.g., PCA or hex values) for plotting the coordinates.
-    :param tar_features: Feature matrix for the target, used to split color values between target and source data.
-    :param shift: Value used to adjust the plot limits around the coordinates for better visualization. Default is 300.
-    :param s: Marker size for the scatter plot points. Default is 0.8.
-    :param boundary_line: Boolean indicating whether to draw boundary lines (horizontal and vertical lines). Default is True.
-    :return: Displays the alignment plot of target, source, and alignment of source coordinates.
+    Optimized plot: target, source, and aligned coordinates with titles.
     """
-    
-    # Create a figure with three subplots, adjusting size and resolution
-    plt.figure(figsize=(10, 3), dpi=300)
-    
-    # First subplot: Plot target coordinates
-    plt.subplot(1, 3, 1)
-    plt.scatter(ad_tar_coor[:, 0], ad_tar_coor[:, 1], marker='o', s=s, c=pca_hex_comb[:len(tar_features.T)])
-    # Set plot limits based on the minimum and maximum target coordinates, with extra padding from 'shift'
-    plt.xlim([ad_tar_coor.min() - shift, ad_tar_coor.max() + shift])
-    plt.ylim([ad_tar_coor.min() - shift, ad_tar_coor.max() + shift])
-    
-    # Second subplot: Plot source coordinates
-    plt.subplot(1, 3, 2)
-    plt.scatter(ad_src_coor[:, 0], ad_src_coor[:, 1], marker='o', s=s, c=pca_hex_comb[len(tar_features.T):])
-    # Ensure consistent plot limits across subplots by using the same limits as the target coordinates
-    plt.xlim([ad_tar_coor.min() - shift, ad_tar_coor.max() + shift])
-    plt.ylim([ad_tar_coor.min() - shift, ad_tar_coor.max() + shift])
-    
-    # Third subplot: Plot alignment of source coordinates
-    plt.subplot(1, 3, 3)
-    plt.scatter(homo_coor[:, 0], homo_coor[:, 1], marker='o', s=s, c=pca_hex_comb[len(tar_features.T):])
-    # Maintain the same plot limits across all subplots for a uniform comparison
-    plt.xlim([ad_tar_coor.min() - shift, ad_tar_coor.max() + shift])
-    plt.ylim([ad_tar_coor.min() - shift, ad_tar_coor.max() + shift])
-    
-    # Optionally draw boundary lines at the minimum x and y values of the target coordinates
-    if boundary_line:
-        plt.axvline(x=ad_tar_coor[:, 0].min(), color='black')  # Vertical boundary line at the minimum x of target coordinates
-        plt.axhline(y=ad_tar_coor[:, 1].min(), color='black')  # Horizontal boundary line at the minimum y of target coordinates
-    
-    # Remove the axis labels and ticks from all subplots for a cleaner appearance
-    plt.axis('off')
-    
-    # Display the plot
+    # Determine common limits
+    coords = np.vstack([ad_tar_coor, ad_src_coor, homo_coor])
+    x_min, x_max = coords[:,0].min() - shift, coords[:,0].max() + shift
+    y_min, y_max = coords[:,1].min() - shift, coords[:,1].max() + shift
+
+    fig, axes = plt.subplots(1, 3, figsize=(10, 3), dpi=150)
+    titles = ["Target ST", "Source ST", "Aligned Source ST"]
+    splits = [len(ad_tar_coor), len(ad_tar_coor)+len(ad_src_coor)]
+
+    for ax, title, data_slice in zip(
+        axes,
+        titles,
+        [(ad_tar_coor, pca_hex_comb[:splits[0]]),
+         (ad_src_coor, pca_hex_comb[splits[0]:splits[1]]),
+         (homo_coor, pca_hex_comb[splits[0]:splits[1]])]
+    ):
+        coords_arr, colors = data_slice
+        ax.scatter(coords_arr[:,0], coords_arr[:,1], s=s, c=colors, marker='o')
+        ax.set_xlim(x_min, x_max)
+        ax.set_ylim(y_min, y_max)
+        ax.set_aspect('equal')
+        if boundary_line:
+            ax.axvline(x=ad_tar_coor[:,0].min(), color='black', linewidth=1)
+            ax.axhline(y=ad_tar_coor[:,1].min(), color='black', linewidth=1)
+        ax.set_title(title)
+        ax.axis('off')
+    plt.tight_layout()
     plt.show()
 
 
 
-def plot_alignment_with_img(ad_tar_coor, ad_src_coor, homo_coor, tar_img, src_img, aligned_image, pca_hex_comb, tar_features):
+def plot_alignment_with_img(
+    ad_tar_coor: np.ndarray,
+    ad_src_coor: np.ndarray,
+    homo_coor: np.ndarray,
+    tar_img,
+    src_img,
+    aligned_image,
+    pca_hex_comb: np.ndarray,
+    tar_features: np.ndarray,
+    s: float = 1.0
+) -> None:
     """
-    Plots the target coordinates and alignment of source coordinates with their respective images in the background.
-
-    :param ad_tar_coor: Numpy array of target coordinates to be plotted in the first and third subplots.
-    :param ad_src_coor: Numpy array of source coordinates to be plotted in the second subplot.
-    :param homo_coor: Numpy array of alignment of source coordinates to be plotted in the third subplot.
-    :param tar_img: Image associated with the target coordinates, used as the background in the first subplot.
-    :param src_img: Image associated with the source coordinates, used as the background in the second subplot.
-    :param aligned_image: Image associated with the aligned coordinates, used as the background in the third subplot.
-    :param pca_hex_comb: Color values (e.g., PCA or hex values) for plotting the coordinates.
-    :param tar_features: Feature matrix for the target, used to split color values between target and source data.
-    :return: Displays the alignment plot of target, source, and alignment of source coordinates with their associated images.
+    Optimized plot with images in the background and subplot titles.
     """
-    
-    # Create a figure with three subplots and set the size and resolution
-    plt.figure(figsize=(10, 8), dpi=150)
-    
-    # First subplot: Plot target coordinates with the target image as the background
-    plt.subplot(1, 3, 1)
-    # Scatter plot for the target coordinates with transparency and small marker size
-    plt.scatter(ad_tar_coor[:, 0], ad_tar_coor[:, 1], marker='o', alpha=0.8, s=1, c=pca_hex_comb[:len(tar_features.T)])
-    # Overlay the target image with some transparency (alpha = 0.3)
-    plt.imshow(tar_img, origin='lower', alpha=0.3)
-    
-    # Second subplot: Plot source coordinates with the source image as the background
-    plt.subplot(1, 3, 2)
-    # Scatter plot for the source coordinates with transparency and small marker size
-    plt.scatter(ad_src_coor[:, 0], ad_src_coor[:, 1], marker='o', alpha=0.8, s=1, c=pca_hex_comb[len(tar_features.T):])
-    # Overlay the source image with some transparency (alpha = 0.3)
-    plt.imshow(src_img, origin='lower', alpha=0.3)
-    
-    # Third subplot: Plot both target and alignment of source coordinates with the aligned image as the background
-    plt.subplot(1, 3, 3)
-    # Scatter plot for the target coordinates with lower opacity (alpha = 0.2)
-    plt.scatter(ad_tar_coor[:, 0], ad_tar_coor[:, 1], marker='o', alpha=0.2, s=1, c=pca_hex_comb[:len(tar_features.T)])
-    # Scatter plot for the homologous coordinates with a '+' marker and the same color mapping
-    plt.scatter(homo_coor[:, 0], homo_coor[:, 1], marker='+', s=1, c=pca_hex_comb[len(tar_features.T):])
-    # Overlay the aligned image with some transparency (alpha = 0.3)
-    plt.imshow(aligned_image, origin='lower', alpha=0.3)
-    
-    # Turn off the axis for all subplots to give a cleaner visual output
-    plt.axis('off')
-    
-    # Display the plots
+    fig, axes = plt.subplots(1, 3, figsize=(15, 5), dpi=150)
+    titles = ["Target + Image", "Source + Image", "Aligned + Image"]
+    splits = [len(tar_features.T), len(tar_features.T) * 2]
+
+    # Data slices for each subplot
+    data_slices = [
+        (ad_tar_coor, pca_hex_comb[:splits[0]], tar_img),
+        (ad_src_coor, pca_hex_comb[splits[0]:splits[1]], src_img),
+        (np.vstack([ad_tar_coor, homo_coor]), 
+         np.concatenate([pca_hex_comb[:splits[0]], pca_hex_comb[splits[0]:splits[1]]]),
+         aligned_image)
+    ]
+
+    for ax, title, (coords_arr, colors, img) in zip(axes, titles, data_slices):
+        ax.imshow(img, origin='lower', alpha=0.3)
+        ax.scatter(coords_arr[:,0], coords_arr[:,1], s=s, c=colors, marker='o')
+        ax.set_aspect('equal')
+        ax.set_title(title)
+        ax.axis('off')
+
+    plt.tight_layout()
     plt.show()
 
 
+def show_image(img, title: str = "Aligned Source Image", origin: str = "lower", cmap=None):
+    """
+    Display a single image with no axes and a title.
+
+    :param img: The image to display (NumPy array, PIL Image, etc.).
+    :param title: Title to show above the image.
+    :param origin: Origin parameter passed to plt.imshow (e.g. 'lower' or 'upper').
+    :param cmap: Optional colormap for grayscale or other single‑channel data.
+    """
+    plt.imshow(img, origin=origin, cmap=cmap)
+    plt.title(title)
+    plt.axis('off')
+    plt.show()
+
 
 def draw_polygon(image, polygon, color='k', thickness=2):
     """
@@ -228,6 +223,9 @@ def plot_heatmap(
     coor,
     similairty,
     image_path=None,
+    polygons=None,
+    polygons_color='k',
+    polygons_thickness=2,
     patch_size=(256, 256),
     save_path=None,
     downsize=32,
@@ -236,9 +234,6 @@ def plot_heatmap(
     boxes=None,
     box_color='k',
     box_thickness=2,
-    polygons=None,
-    polygons_color='k',
-    polygons_thickness=2,
     image_alpha=0.5
 ):
     """
@@ -316,7 +311,7 @@ def plot_heatmap(
 
 
 
-def show_images_side_by_side(image1, image2, title1=None, title2=None):
+def show_images_side_by_side(image1, image2, title1='Annotated H&E Image', title2='Similatrity Heatmap'):
     """
     Displays two images side by side in a single figure.
 
@@ -328,7 +323,7 @@ def show_images_side_by_side(image1, image2, title1=None, title2=None):
     """
     
     # Create a figure with 2 subplots (1 row, 2 columns), and set the figure size
-    fig, ax = plt.subplots(1, 2, figsize=(15,8))
+    fig, ax = plt.subplots(1, 2, figsize=(8,6), dpi=150)
     
     # Display the first image on the first subplot
     ax[0].imshow(image1)
@@ -364,7 +359,7 @@ def plot_img_with_annotation(fullres_img, roi_polygon, linewidth, xlim, ylim):
     """
     
     # Create a new figure with a fixed size for displaying the image and annotations
-    plt.figure(figsize=(10, 10))
+    plt.figure(figsize=(12, 12), dpi=150)
     
     # Display the full-resolution image
     plt.imshow(fullres_img)
@@ -403,7 +398,7 @@ def plot_annotation_heatmap(st_ad, roi_polygon, s, linewidth, xlim, ylim):
     """
     
     # Create a new figure with a fixed size for displaying the heatmap and annotations
-    plt.figure(figsize=(10, 10))
+    plt.figure(figsize=(12, 12), dpi=150)
     
     # Scatter plot for the spatial transcriptomics data.
     # The 'spatial' coordinates are plotted with color intensity based on 'bulk_simi' values.

src/build/lib/loki/utils.py

@@ -11,175 +11,107 @@ from open_clip import create_model_from_pretrained, get_tokenizer
 
 
 
-def load_model(model_path, device):
+import os
+from typing import List, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+from PIL import Image
+import pandas as pd
+
+# --- Model loading --------------------------------------------------------
+
+def load_model(
+    model_path: str,
+    device: Union[str, torch.device]
+) -> Tuple[torch.nn.Module, callable, callable]:
     """
-    Loads a pretrained CoCa (CLIP-like) model, along with its preprocessing function and tokenizer, 
-    using the specified model checkpoint.
-
-    :param model_path: File path or URL to the pretrained model checkpoint. This is passed to 
-                       `create_model_from_pretrained` as the `pretrained` argument.
-    :type model_path: str
-    :param device: The device on which to load the model (e.g., 'cpu' or 'cuda').
-    :type device: str or torch.device
-    :return: A tuple `(model, preprocess, tokenizer)` where:
-             - model: The loaded CoCa model.
-             - preprocess: A function or transform that preprocesses input data for the model.
-             - tokenizer: A tokenizer appropriate for textual input to the model.
-    :rtype: (nn.Module, callable, callable)
+    Load pretrained OmiCLIP (COCA ViT‑L‑14) model, its image preprocess, and tokenizer.
     """
-    # Create the model and its preprocessing transform from the specified checkpoint
     model, preprocess = create_model_from_pretrained(
         "coca_ViT-L-14", device=device, pretrained=model_path
     )
-    
-    # Retrieve a tokenizer compatible with the "coca_ViT-L-14" architecture
-    tokenizer = get_tokenizer('coca_ViT-L-14')
-
+    tokenizer = get_tokenizer("coca_ViT-L-14")
+    model.to(device).eval()
     return model, preprocess, tokenizer
 
+# --- Image encoding -------------------------------------------------------
 
-
-def encode_image(model, preprocess, image):
-    """
-    Encodes an image into a normalized feature embedding using the specified model and preprocessing function.
-
-    :param model: A model object that provides an `encode_image` method (e.g., a CLIP or CoCa model).
-    :type model: torch.nn.Module
-    :param preprocess: A preprocessing function that transforms the input image into a tensor 
-                       suitable for the model. Typically something returning a PyTorch tensor.
-    :type preprocess: callable
-    :param image: The input image (PIL Image, NumPy array, or other format supported by `preprocess`).
-    :type image: PIL.Image.Image or numpy.ndarray
-    :return: A single normalized image embedding as a PyTorch tensor of shape (1, embedding_dim).
-    :rtype: torch.Tensor
+def encode_images(
+    model: torch.nn.Module,
+    preprocess: callable,
+    image_paths: List[str],
+    device: Union[str, torch.device]
+) -> torch.Tensor:
     """
-    # Preprocess the image, then stack to create a batch of size 1
-    image_input = torch.stack([preprocess(image)])
-
-    # Generate the image features without gradient tracking
-    with torch.no_grad():
-        image_features = model.encode_image(image_input)
-
-    # Normalize embeddings across the feature dimension (L2 normalization)
-    image_embeddings = F.normalize(image_features, p=2, dim=-1)
-
-    return image_embeddings
-
-
-
-def encode_image_patches(model, preprocess, data_dir, img_list):
+    Batch–encode a list of image file paths into L2‑normalized embeddings.
+    Returns a tensor of shape (N, D).
     """
-    Encodes multiple image patches into normalized feature embeddings using a specified model and preprocess function.
+    # Load & preprocess all images
+    imgs = [preprocess(Image.open(p)) for p in image_paths]
+    batch = torch.stack(imgs, dim=0).to(device)           # (N, C, H, W)
     
-    :param model: A model object that provides an `encode_image` method (e.g., a CLIP or CoCa model).
-    :type model: torch.nn.Module
-    :param preprocess: A preprocessing function that transforms the input image into a tensor 
-                       suitable for the model. Typically something returning a PyTorch tensor.
-    :type preprocess: callable
-    :param data_dir: The base directory containing image data.
-    :type data_dir: str
-    :param img_list: A list of image filenames (strings). Each filename corresponds to a patch image 
-                     stored in `data_dir/demo_data/patch/`.
-    :type img_list: list[str]
-    :return: A PyTorch tensor of shape (N, 1, embedding_dim), containing the normalized embeddings 
-             for each image in `img_list`.
-    :rtype: torch.Tensor
-    """
+    with torch.no_grad():
+        feats = model.encode_image(batch)                 # (N, D)
+    return F.normalize(feats, p=2, dim=-1)                # (N, D)
 
-    # Prepare a list to hold each image's feature embedding
-    image_embeddings = []
 
-    # Loop through each image name in the provided list
-    for img_name in img_list:
-        # Build the path to the patch image and open it
-        image_path = os.path.join(data_dir, 'demo_data', 'patch', img_name)
-        image = Image.open(image_path)
+    # # Loop through each image name in the provided list
+    # for img_name in img_list:
+    #     # Build the path to the patch image and open it
+    #     image_path = os.path.join(data_dir, 'demo_data', 'patch', img_name)
+    #     image = Image.open(image_path)
 
-        # Encode the image using the model & preprocess; returns shape (1, embedding_dim)
-        image_features = encode_image(model, preprocess, image)
+    #     # Encode the image using the model & preprocess; returns shape (1, embedding_dim)
+    #     image_features = encode_image(model, preprocess, image)
 
-        # Accumulate the feature embeddings in the list
-        image_embeddings.append(image_features)
+    #     # Accumulate the feature embeddings in the list
+    #     image_embeddings.append(image_features)
 
-    # Convert the list of embeddings to a NumPy array, then to a PyTorch tensor
-    # Resulting shape will be (N, 1, embedding_dim)
-    image_embeddings = torch.from_numpy(np.array(image_embeddings))
+    # # Convert the list of embeddings to a NumPy array, then to a PyTorch tensor
+    # # Resulting shape will be (N, 1, embedding_dim)
+    # image_embeddings = torch.from_numpy(np.array(image_embeddings))
 
-    # Normalize all embeddings across the feature dimension (L2 normalization)
-    image_embeddings = F.normalize(image_embeddings, p=2, dim=-1)
+    # # Normalize all embeddings across the feature dimension (L2 normalization)
+    # image_embeddings = F.normalize(image_embeddings, p=2, dim=-1)
 
-    return image_embeddings
+    # return image_embeddings
 
 
+# --- Text encoding --------------------------------------------------------
 
-def encode_text(model, tokenizer, text):
+def encode_texts(
+    model: torch.nn.Module,
+    tokenizer: callable,
+    texts: List[str],
+    device: Union[str, torch.device]
+) -> torch.Tensor:
     """
-    Encodes text into a normalized feature embedding using a specified model and tokenizer.
-
-    :param model: A model object that provides an `encode_text` method (e.g., a CLIP-like or CoCa model).
-    :type model: torch.nn.Module
-    :param tokenizer: A tokenizer function that converts the input text into a format suitable for `model.encode_text`.
-                      Typically returns token IDs, attention masks, etc. as a torch.Tensor or similar structure.
-    :type tokenizer: callable
-    :param text: The input text (string or list of strings) to be encoded.
-    :type text: str or list[str]
-    :return: A PyTorch tensor of shape (batch_size, embedding_dim) containing the L2-normalized text embeddings.
-    :rtype: torch.Tensor
+    Batch–encode a list of strings into L2‑normalized embeddings.
+    Returns a tensor of shape (N, D).
     """
-
-    # Convert text to the appropriate tokenized representation
-    text_input = tokenizer(text)
-
-    # Run the model in no-grad mode (not tracking gradients, saving memory and compute)
+    # Tokenizer returns a dict of tensors
+    text_inputs = tokenizer(texts)
+    
     with torch.no_grad():
-        text_features = model.encode_text(text_input)
-
-    # Normalize embeddings to unit length
-    text_embeddings = F.normalize(text_features, p=2, dim=-1)
+        feats = model.encode_text(text_inputs)             # (N, D)
+    return F.normalize(feats, p=2, dim=-1)                # (N, D)
 
-    return text_embeddings
 
-
-
-def encode_text_df(model, tokenizer, df, col_name):
+def encode_text_df(
+    model: torch.nn.Module,
+    tokenizer: callable,
+    df: pd.DataFrame,
+    col_name: str,
+    device: Union[str, torch.device]
+) -> torch.Tensor:
     """
-    Encodes text from a specified column in a pandas DataFrame using the given model and tokenizer,
-    returning a PyTorch tensor of normalized text embeddings.
-
-    :param model: A model object that provides an `encode_text` method (e.g., a CLIP-like or CoCa model).
-    :type model: torch.nn.Module
-    :param tokenizer: A tokenizer function that converts the input text into a format suitable for `model.encode_text`.
-    :type tokenizer: callable
-    :param df: A pandas DataFrame from which text will be extracted.
-    :type df: pandas.DataFrame
-    :param col_name: The name of the column in `df` that contains the text to be encoded.
-    :type col_name: str
-    :return: A PyTorch tensor containing the L2-normalized text embeddings, 
-             where the shape is (number_of_rows, embedding_dim).
-    :rtype: torch.Tensor
+    Encodes an entire DataFrame column into (N, D) embeddings.
     """
+    texts = df[col_name].astype(str).tolist()
+    return encode_texts(model, tokenizer, texts, device)
 
-    # Prepare a list to hold each row's text embedding
-    text_embeddings = []
-
-    # Loop through each index in the DataFrame
-    for idx in df.index:
-        # Retrieve text from the specified column for the current row
-        text = df[df.index == idx][col_name][0]
-
-        # Encode the text using the provided model and tokenizer
-        text_features = encode_text(model, tokenizer, text)
-
-        # Accumulate the embedding tensor
-        text_embeddings.append(text_features)
-
-    # Convert the list of embeddings (likely shape [N, embedding_dim]) into a NumPy array, then to a torch tensor
-    text_embeddings = torch.from_numpy(np.array(text_embeddings))
-
-    # Normalize embeddings to unit length across the feature dimension
-    text_embeddings = F.normalize(text_embeddings, p=2, dim=-1)
-
-    return text_embeddings
 
 
 

src/loki.egg-info/PKG-INFO

@@ -1,13 +1,14 @@
 Metadata-Version: 2.1
 Name: loki
 Version: 0.0.1
-Summary: The Loki platform offers 5 core functions: tissue alignment, cell type decomposition, tissue annotation, image-transcriptomics retrieval, and ST gene expression prediction
+Summary: The Loki platform offers 5 core functions: tissue alignment, tissue annotation, cell type decomposition, image-transcriptomics retrieval, and ST gene expression prediction
 Author: Weiqing Chen
 Author-email: wec4005@med.cornell.edu
 Classifier: Programming Language :: Python :: 3
-Classifier: License :: OSI Approved :: MIT License
+Classifier: License :: BSD 3-Clause License
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.9
+License-File: LICENSE
 Requires-Dist: anndata==0.10.9
 Requires-Dist: matplotlib==3.9.2
 Requires-Dist: numpy==1.25.0

src/loki.egg-info/SOURCES.txt

@@ -1,4 +1,4 @@
-README.md
+LICENSE
 setup.py
 loki/__init__.py
 loki/align.py

src/loki/plot.py

@@ -8,107 +8,102 @@ import numpy as np
 from tqdm import tqdm
 
 
-
-def plot_alignment(ad_tar_coor, ad_src_coor, homo_coor, pca_hex_comb, tar_features, shift=300, s=0.8, boundary_line=True):
+def plot_alignment(
+    ad_tar_coor: np.ndarray,
+    ad_src_coor: np.ndarray,
+    homo_coor: np.ndarray,
+    pca_hex_comb: np.ndarray,
+    tar_features: np.ndarray,
+    shift: float = 300,
+    s: float = 0.8,
+    boundary_line: bool = True
+) -> None:
     """
-    Plots the target coordinates and alignment of source coordinates.
-
-    :param ad_tar_coor: Numpy array of target coordinates to be plotted in the first subplot.
-    :param ad_src_coor: Numpy array of source coordinates to be plotted in the second subplot.
-    :param homo_coor: Numpy array of alignment of source coordinates to be plotted in the third subplot.
-    :param pca_hex_comb: Color values (e.g., PCA or hex values) for plotting the coordinates.
-    :param tar_features: Feature matrix for the target, used to split color values between target and source data.
-    :param shift: Value used to adjust the plot limits around the coordinates for better visualization. Default is 300.
-    :param s: Marker size for the scatter plot points. Default is 0.8.
-    :param boundary_line: Boolean indicating whether to draw boundary lines (horizontal and vertical lines). Default is True.
-    :return: Displays the alignment plot of target, source, and alignment of source coordinates.
+    Optimized plot: target, source, and aligned coordinates with titles.
     """
-    
-    # Create a figure with three subplots, adjusting size and resolution
-    plt.figure(figsize=(10, 3), dpi=300)
-    
-    # First subplot: Plot target coordinates
-    plt.subplot(1, 3, 1)
-    plt.scatter(ad_tar_coor[:, 0], ad_tar_coor[:, 1], marker='o', s=s, c=pca_hex_comb[:len(tar_features.T)])
-    # Set plot limits based on the minimum and maximum target coordinates, with extra padding from 'shift'
-    plt.xlim([ad_tar_coor.min() - shift, ad_tar_coor.max() + shift])
-    plt.ylim([ad_tar_coor.min() - shift, ad_tar_coor.max() + shift])
-    
-    # Second subplot: Plot source coordinates
-    plt.subplot(1, 3, 2)
-    plt.scatter(ad_src_coor[:, 0], ad_src_coor[:, 1], marker='o', s=s, c=pca_hex_comb[len(tar_features.T):])
-    # Ensure consistent plot limits across subplots by using the same limits as the target coordinates
-    plt.xlim([ad_tar_coor.min() - shift, ad_tar_coor.max() + shift])
-    plt.ylim([ad_tar_coor.min() - shift, ad_tar_coor.max() + shift])
-    
-    # Third subplot: Plot alignment of source coordinates
-    plt.subplot(1, 3, 3)
-    plt.scatter(homo_coor[:, 0], homo_coor[:, 1], marker='o', s=s, c=pca_hex_comb[len(tar_features.T):])
-    # Maintain the same plot limits across all subplots for a uniform comparison
-    plt.xlim([ad_tar_coor.min() - shift, ad_tar_coor.max() + shift])
-    plt.ylim([ad_tar_coor.min() - shift, ad_tar_coor.max() + shift])
-    
-    # Optionally draw boundary lines at the minimum x and y values of the target coordinates
-    if boundary_line:
-        plt.axvline(x=ad_tar_coor[:, 0].min(), color='black')  # Vertical boundary line at the minimum x of target coordinates
-        plt.axhline(y=ad_tar_coor[:, 1].min(), color='black')  # Horizontal boundary line at the minimum y of target coordinates
-    
-    # Remove the axis labels and ticks from all subplots for a cleaner appearance
-    plt.axis('off')
-    
-    # Display the plot
+    # Determine common limits
+    coords = np.vstack([ad_tar_coor, ad_src_coor, homo_coor])
+    x_min, x_max = coords[:,0].min() - shift, coords[:,0].max() + shift
+    y_min, y_max = coords[:,1].min() - shift, coords[:,1].max() + shift
+
+    fig, axes = plt.subplots(1, 3, figsize=(10, 3), dpi=150)
+    titles = ["Target ST", "Source ST", "Aligned Source ST"]
+    splits = [len(ad_tar_coor), len(ad_tar_coor)+len(ad_src_coor)]
+
+    for ax, title, data_slice in zip(
+        axes,
+        titles,
+        [(ad_tar_coor, pca_hex_comb[:splits[0]]),
+         (ad_src_coor, pca_hex_comb[splits[0]:splits[1]]),
+         (homo_coor, pca_hex_comb[splits[0]:splits[1]])]
+    ):
+        coords_arr, colors = data_slice
+        ax.scatter(coords_arr[:,0], coords_arr[:,1], s=s, c=colors, marker='o')
+        ax.set_xlim(x_min, x_max)
+        ax.set_ylim(y_min, y_max)
+        ax.set_aspect('equal')
+        if boundary_line:
+            ax.axvline(x=ad_tar_coor[:,0].min(), color='black', linewidth=1)
+            ax.axhline(y=ad_tar_coor[:,1].min(), color='black', linewidth=1)
+        ax.set_title(title)
+        ax.axis('off')
+    plt.tight_layout()
     plt.show()
 
 
 
-def plot_alignment_with_img(ad_tar_coor, ad_src_coor, homo_coor, tar_img, src_img, aligned_image, pca_hex_comb, tar_features):
+def plot_alignment_with_img(
+    ad_tar_coor: np.ndarray,
+    ad_src_coor: np.ndarray,
+    homo_coor: np.ndarray,
+    tar_img,
+    src_img,
+    aligned_image,
+    pca_hex_comb: np.ndarray,
+    tar_features: np.ndarray,
+    s: float = 1.0
+) -> None:
     """
-    Plots the target coordinates and alignment of source coordinates with their respective images in the background.
-
-    :param ad_tar_coor: Numpy array of target coordinates to be plotted in the first and third subplots.
-    :param ad_src_coor: Numpy array of source coordinates to be plotted in the second subplot.
-    :param homo_coor: Numpy array of alignment of source coordinates to be plotted in the third subplot.
-    :param tar_img: Image associated with the target coordinates, used as the background in the first subplot.
-    :param src_img: Image associated with the source coordinates, used as the background in the second subplot.
-    :param aligned_image: Image associated with the aligned coordinates, used as the background in the third subplot.
-    :param pca_hex_comb: Color values (e.g., PCA or hex values) for plotting the coordinates.
-    :param tar_features: Feature matrix for the target, used to split color values between target and source data.
-    :return: Displays the alignment plot of target, source, and alignment of source coordinates with their associated images.
+    Optimized plot with images in the background and subplot titles.
     """
-    
-    # Create a figure with three subplots and set the size and resolution
-    plt.figure(figsize=(10, 8), dpi=150)
-    
-    # First subplot: Plot target coordinates with the target image as the background
-    plt.subplot(1, 3, 1)
-    # Scatter plot for the target coordinates with transparency and small marker size
-    plt.scatter(ad_tar_coor[:, 0], ad_tar_coor[:, 1], marker='o', alpha=0.8, s=1, c=pca_hex_comb[:len(tar_features.T)])
-    # Overlay the target image with some transparency (alpha = 0.3)
-    plt.imshow(tar_img, origin='lower', alpha=0.3)
-    
-    # Second subplot: Plot source coordinates with the source image as the background
-    plt.subplot(1, 3, 2)
-    # Scatter plot for the source coordinates with transparency and small marker size
-    plt.scatter(ad_src_coor[:, 0], ad_src_coor[:, 1], marker='o', alpha=0.8, s=1, c=pca_hex_comb[len(tar_features.T):])
-    # Overlay the source image with some transparency (alpha = 0.3)
-    plt.imshow(src_img, origin='lower', alpha=0.3)
-    
-    # Third subplot: Plot both target and alignment of source coordinates with the aligned image as the background
-    plt.subplot(1, 3, 3)
-    # Scatter plot for the target coordinates with lower opacity (alpha = 0.2)
-    plt.scatter(ad_tar_coor[:, 0], ad_tar_coor[:, 1], marker='o', alpha=0.2, s=1, c=pca_hex_comb[:len(tar_features.T)])
-    # Scatter plot for the homologous coordinates with a '+' marker and the same color mapping
-    plt.scatter(homo_coor[:, 0], homo_coor[:, 1], marker='+', s=1, c=pca_hex_comb[len(tar_features.T):])
-    # Overlay the aligned image with some transparency (alpha = 0.3)
-    plt.imshow(aligned_image, origin='lower', alpha=0.3)
-    
-    # Turn off the axis for all subplots to give a cleaner visual output
-    plt.axis('off')
-    
-    # Display the plots
+    fig, axes = plt.subplots(1, 3, figsize=(15, 5), dpi=150)
+    titles = ["Target + Image", "Source + Image", "Aligned + Image"]
+    splits = [len(tar_features.T), len(tar_features.T) * 2]
+
+    # Data slices for each subplot
+    data_slices = [
+        (ad_tar_coor, pca_hex_comb[:splits[0]], tar_img),
+        (ad_src_coor, pca_hex_comb[splits[0]:splits[1]], src_img),
+        (np.vstack([ad_tar_coor, homo_coor]), 
+         np.concatenate([pca_hex_comb[:splits[0]], pca_hex_comb[splits[0]:splits[1]]]),
+         aligned_image)
+    ]
+
+    for ax, title, (coords_arr, colors, img) in zip(axes, titles, data_slices):
+        ax.imshow(img, origin='lower', alpha=0.3)
+        ax.scatter(coords_arr[:,0], coords_arr[:,1], s=s, c=colors, marker='o')
+        ax.set_aspect('equal')
+        ax.set_title(title)
+        ax.axis('off')
+
+    plt.tight_layout()
     plt.show()
 
 
+def show_image(img, title: str = "Aligned Source Image", origin: str = "lower", cmap=None):
+    """
+    Display a single image with no axes and a title.
+
+    :param img: The image to display (NumPy array, PIL Image, etc.).
+    :param title: Title to show above the image.
+    :param origin: Origin parameter passed to plt.imshow (e.g. 'lower' or 'upper').
+    :param cmap: Optional colormap for grayscale or other single‑channel data.
+    """
+    plt.imshow(img, origin=origin, cmap=cmap)
+    plt.title(title)
+    plt.axis('off')
+    plt.show()
+
 
 def draw_polygon(image, polygon, color='k', thickness=2):
     """
@@ -228,6 +223,9 @@ def plot_heatmap(
     coor,
     similairty,
     image_path=None,
+    polygons=None,
+    polygons_color='k',
+    polygons_thickness=2,
     patch_size=(256, 256),
     save_path=None,
     downsize=32,
@@ -236,9 +234,6 @@ def plot_heatmap(
     boxes=None,
     box_color='k',
     box_thickness=2,
-    polygons=None,
-    polygons_color='k',
-    polygons_thickness=2,
     image_alpha=0.5
 ):
     """
@@ -316,7 +311,7 @@ def plot_heatmap(
 
 
 
-def show_images_side_by_side(image1, image2, title1=None, title2=None):
+def show_images_side_by_side(image1, image2, title1='Annotated H&E Image', title2='Similatrity Heatmap'):
     """
     Displays two images side by side in a single figure.
 
@@ -328,7 +323,7 @@ def show_images_side_by_side(image1, image2, title1=None, title2=None):
     """
     
     # Create a figure with 2 subplots (1 row, 2 columns), and set the figure size
-    fig, ax = plt.subplots(1, 2, figsize=(15,8))
+    fig, ax = plt.subplots(1, 2, figsize=(8,6), dpi=150)
     
     # Display the first image on the first subplot
     ax[0].imshow(image1)
@@ -364,7 +359,7 @@ def plot_img_with_annotation(fullres_img, roi_polygon, linewidth, xlim, ylim):
     """
     
     # Create a new figure with a fixed size for displaying the image and annotations
-    plt.figure(figsize=(10, 10))
+    plt.figure(figsize=(12, 12), dpi=150)
     
     # Display the full-resolution image
     plt.imshow(fullres_img)
@@ -403,7 +398,7 @@ def plot_annotation_heatmap(st_ad, roi_polygon, s, linewidth, xlim, ylim):
     """
     
     # Create a new figure with a fixed size for displaying the heatmap and annotations
-    plt.figure(figsize=(10, 10))
+    plt.figure(figsize=(12, 12), dpi=150)
     
     # Scatter plot for the spatial transcriptomics data.
     # The 'spatial' coordinates are plotted with color intensity based on 'bulk_simi' values.

src/loki/utils.py

@@ -11,175 +11,107 @@ from open_clip import create_model_from_pretrained, get_tokenizer
 
 
 
-def load_model(model_path, device):
+import os
+from typing import List, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+from PIL import Image
+import pandas as pd
+
+# --- Model loading --------------------------------------------------------
+
+def load_model(
+    model_path: str,
+    device: Union[str, torch.device]
+) -> Tuple[torch.nn.Module, callable, callable]:
     """
-    Loads a pretrained OmiCLIP model, along with its preprocessing function and tokenizer, 
-    using the specified model checkpoint.
-
-    :param model_path: File path to the pretrained model checkpoint. This is passed to 
-                       `create_model_from_pretrained` as the `pretrained` argument.
-    :type model_path: str
-    :param device: The device on which to load the model (e.g., 'cpu' or 'cuda').
-    :type device: str or torch.device
-    :return: A tuple `(model, preprocess, tokenizer)` where:
-             - model: The loaded OmiCLIP model.
-             - preprocess: A function or transform that preprocesses input data for the model.
-             - tokenizer: A tokenizer appropriate for textual input to the model.
-    :rtype: (nn.Module, callable, callable)
+    Load pretrained OmiCLIP (COCA ViT‑L‑14) model, its image preprocess, and tokenizer.
     """
-    # Create the model and its preprocessing transform from the specified checkpoint
     model, preprocess = create_model_from_pretrained(
         "coca_ViT-L-14", device=device, pretrained=model_path
     )
-    
-    # Retrieve a tokenizer compatible with the "coca_ViT-L-14" architecture
-    tokenizer = get_tokenizer('coca_ViT-L-14')
-
+    tokenizer = get_tokenizer("coca_ViT-L-14")
+    model.to(device).eval()
     return model, preprocess, tokenizer
 
+# --- Image encoding -------------------------------------------------------
 
-
-def encode_image(model, preprocess, image):
-    """
-    Encodes an image into a normalized feature embedding using the specified model and preprocessing function.
-
-    :param model: A model object that provides an `encode_image` method.
-    :type model: torch.nn.Module
-    :param preprocess: A preprocessing function that transforms the input image into a tensor 
-                       suitable for the model. Typically something returning a PyTorch tensor.
-    :type preprocess: callable
-    :param image: The input image (PIL Image, NumPy array, or other format supported by `preprocess`).
-    :type image: PIL.Image.Image or numpy.ndarray
-    :return: A single normalized image embedding as a PyTorch tensor of shape (1, embedding_dim).
-    :rtype: torch.Tensor
+def encode_images(
+    model: torch.nn.Module,
+    preprocess: callable,
+    image_paths: List[str],
+    device: Union[str, torch.device]
+) -> torch.Tensor:
     """
-    # Preprocess the image, then stack to create a batch of size 1
-    image_input = torch.stack([preprocess(image)])
-
-    # Generate the image features without gradient tracking
-    with torch.no_grad():
-        image_features = model.encode_image(image_input)
-
-    # Normalize embeddings across the feature dimension (L2 normalization)
-    image_embeddings = F.normalize(image_features, p=2, dim=-1)
-
-    return image_embeddings
-
-
-
-def encode_image_patches(model, preprocess, data_dir, img_list):
+    Batch–encode a list of image file paths into L2‑normalized embeddings.
+    Returns a tensor of shape (N, D).
     """
-    Encodes multiple image patches into normalized feature embeddings using a specified model and preprocess function.
+    # Load & preprocess all images
+    imgs = [preprocess(Image.open(p)) for p in image_paths]
+    batch = torch.stack(imgs, dim=0).to(device)           # (N, C, H, W)
     
-    :param model: A model object that provides an `encode_image` method.
-    :type model: torch.nn.Module
-    :param preprocess: A preprocessing function that transforms the input image into a tensor 
-                       suitable for the model. Typically something returning a PyTorch tensor.
-    :type preprocess: callable
-    :param data_dir: The base directory containing image data.
-    :type data_dir: str
-    :param img_list: A list of image filenames (strings). Each filename corresponds to a patch image 
-                     stored in `data_dir/demo_data/patch/`.
-    :type img_list: list[str]
-    :return: A PyTorch tensor of shape (N, 1, embedding_dim), containing the normalized embeddings 
-             for each image in `img_list`.
-    :rtype: torch.Tensor
-    """
+    with torch.no_grad():
+        feats = model.encode_image(batch)                 # (N, D)
+    return F.normalize(feats, p=2, dim=-1)                # (N, D)
 
-    # Prepare a list to hold each image's feature embedding
-    image_embeddings = []
 
-    # Loop through each image name in the provided list
-    for img_name in img_list:
-        # Build the path to the patch image and open it
-        image_path = os.path.join(data_dir, 'demo_data', 'patch', img_name)
-        image = Image.open(image_path)
+    # # Loop through each image name in the provided list
+    # for img_name in img_list:
+    #     # Build the path to the patch image and open it
+    #     image_path = os.path.join(data_dir, 'demo_data', 'patch', img_name)
+    #     image = Image.open(image_path)
 
-        # Encode the image using the model & preprocess; returns shape (1, embedding_dim)
-        image_features = encode_image(model, preprocess, image)
+    #     # Encode the image using the model & preprocess; returns shape (1, embedding_dim)
+    #     image_features = encode_image(model, preprocess, image)
 
-        # Accumulate the feature embeddings in the list
-        image_embeddings.append(image_features)
+    #     # Accumulate the feature embeddings in the list
+    #     image_embeddings.append(image_features)
 
-    # Convert the list of embeddings to a NumPy array, then to a PyTorch tensor
-    # Resulting shape will be (N, 1, embedding_dim)
-    image_embeddings = torch.from_numpy(np.array(image_embeddings))
+    # # Convert the list of embeddings to a NumPy array, then to a PyTorch tensor
+    # # Resulting shape will be (N, 1, embedding_dim)
+    # image_embeddings = torch.from_numpy(np.array(image_embeddings))
 
-    # Normalize all embeddings across the feature dimension (L2 normalization)
-    image_embeddings = F.normalize(image_embeddings, p=2, dim=-1)
+    # # Normalize all embeddings across the feature dimension (L2 normalization)
+    # image_embeddings = F.normalize(image_embeddings, p=2, dim=-1)
 
-    return image_embeddings
+    # return image_embeddings
 
 
+# --- Text encoding --------------------------------------------------------
 
-def encode_text(model, tokenizer, text):
+def encode_texts(
+    model: torch.nn.Module,
+    tokenizer: callable,
+    texts: List[str],
+    device: Union[str, torch.device]
+) -> torch.Tensor:
     """
-    Encodes text into a normalized feature embedding using a specified model and tokenizer.
-
-    :param model: A model object that provides an `encode_text` method.
-    :type model: torch.nn.Module
-    :param tokenizer: A tokenizer function that converts the input text into a format suitable for `model.encode_text`.
-                      Typically returns token IDs, attention masks, etc. as a torch.Tensor or similar structure.
-    :type tokenizer: callable
-    :param text: The input text (string or list of strings) to be encoded.
-    :type text: str or list[str]
-    :return: A PyTorch tensor of shape (batch_size, embedding_dim) containing the L2-normalized text embeddings.
-    :rtype: torch.Tensor
+    Batch–encode a list of strings into L2‑normalized embeddings.
+    Returns a tensor of shape (N, D).
     """
-
-    # Convert text to the appropriate tokenized representation
-    text_input = tokenizer(text)
-
-    # Run the model in no-grad mode (not tracking gradients, saving memory and compute)
+    # Tokenizer returns a dict of tensors
+    text_inputs = tokenizer(texts)
+    
     with torch.no_grad():
-        text_features = model.encode_text(text_input)
-
-    # Normalize embeddings to unit length
-    text_embeddings = F.normalize(text_features, p=2, dim=-1)
+        feats = model.encode_text(text_inputs)             # (N, D)
+    return F.normalize(feats, p=2, dim=-1)                # (N, D)
 
-    return text_embeddings
 
-
-
-def encode_text_df(model, tokenizer, df, col_name):
+def encode_text_df(
+    model: torch.nn.Module,
+    tokenizer: callable,
+    df: pd.DataFrame,
+    col_name: str,
+    device: Union[str, torch.device]
+) -> torch.Tensor:
     """
-    Encodes text from a specified column in a pandas DataFrame using the given model and tokenizer,
-    returning a PyTorch tensor of normalized text embeddings.
-
-    :param model: A model object that provides an `encode_text` method.
-    :type model: torch.nn.Module
-    :param tokenizer: A tokenizer function that converts the input text into a format suitable for `model.encode_text`.
-    :type tokenizer: callable
-    :param df: A pandas DataFrame from which text will be extracted.
-    :type df: pandas.DataFrame
-    :param col_name: The name of the column in `df` that contains the text to be encoded.
-    :type col_name: str
-    :return: A PyTorch tensor containing the L2-normalized text embeddings, 
-             where the shape is (number_of_rows, embedding_dim).
-    :rtype: torch.Tensor
+    Encodes an entire DataFrame column into (N, D) embeddings.
     """
+    texts = df[col_name].astype(str).tolist()
+    return encode_texts(model, tokenizer, texts, device)
 
-    # Prepare a list to hold each row's text embedding
-    text_embeddings = []
-
-    # Loop through each index in the DataFrame
-    for idx in df.index:
-        # Retrieve text from the specified column for the current row
-        text = df[df.index == idx][col_name][0]
-
-        # Encode the text using the provided model and tokenizer
-        text_features = encode_text(model, tokenizer, text)
-
-        # Accumulate the embedding tensor
-        text_embeddings.append(text_features)
-
-    # Convert the list of embeddings (likely shape [N, embedding_dim]) into a NumPy array, then to a torch tensor
-    text_embeddings = torch.from_numpy(np.array(text_embeddings))
-
-    # Normalize embeddings to unit length across the feature dimension
-    text_embeddings = F.normalize(text_embeddings, p=2, dim=-1)
-
-    return text_embeddings
 
 
 

src/requirements.txt

@@ -11,4 +11,5 @@ torchvision==0.18.1
 open_clip_torch==2.26.1
 pillow==10.4.0
 ipykernel==6.29.5
+ipywidgets==8.1.6