Source code for plot

import numpy as np
import matplotlib.pyplot as plt
import os
from tqdm import tqdm
import cv2 as cv
import matplotlib.patches as patches
import pandas as pd
from typing import List, Dict, Tuple, Optional, Union


[docs] class Plot_Visium: def __init__( self, segmentation, boundary_dict: Dict, type_list: List[str], colors: Optional[np.ndarray] = None, ): """ Args: segmentation: Object of spotiphy.segmentation.Segmentation. boundary_dict: Output of function spotiphy.segmentation.cell_boundary. type_list: List of the cell types. colors: Colors of each cell type. """ self.img = segmentation.img # Original image. if np.max(self.img) <= 1: self.img = (self.img * 255).astype(np.int32) self.img_seg = segmentation.label # Segmented image. self.n_cell_df = ( segmentation.n_cell_df ) # Dataframe shown the cells in each spot. self.nucleus_boundary = segmentation.nucleus_boundary.astype( np.int32 ) # Segmented nucleus boundaries. self.nucleus_df = segmentation.nucleus_df self.spot_radius = segmentation.spot_radius self.spot_centers = segmentation.spot_center self.nucleus_centers = self.nucleus_df[["x", "y"]].values.astype( np.int32 ) # Segmented nucleus centers. self.img_cell = boundary_dict["img_cell"].astype( np.int32 ) # Inferred cell location. self.type_mask = None self.img_boundary = boundary_dict["cell_boundary"] self.individual_boundary = boundary_dict["individual_boundary"] self.img_size = self.img_seg.shape[:2] self.type_list = type_list if colors is None: self.colors = np.array( list(plt.get_cmap("tab20b").colors) + list(plt.get_cmap("tab20c").colors) ) self.colors = (self.colors * 255).astype(np.int32) else: self.colors = colors self.color_dict = { type_: list(self.colors[i]) for i, type_ in enumerate(type_list) }
[docs] def plot( self, background: bool = False, cell: str = "both", shape: str = "cell", circle_size: int = 10, boundary: Optional[str] = None, save: Optional[str] = "Visium_plot.png", background_alpha: float = 0.5, spot: bool = True, spot_width: int = 2, spot_color: Tuple[int, int, int] = (0, 0, 255), cell_boundary_color: Tuple[int, int, int] = (100, 100, 100), dpi: int = 300, ) -> None: """ Args: background: If show the background. cell: Which group of cell shapes to plot? [both, in, out] shape: Which shape to plot? [cell, nucleus, circle] circle_size: Size of the nuclei. boundary: Which group of cell boundary to plot? [both, in, out] save: If not none, save the figure to the path. background_alpha: Opacity of the background figure. spot: If plot the spot. spot_width: Width of the spot. spot_color: Color of the spot. cell_boundary_color: Color of the cell boundaries. dpi: DPI of the image plotted. """ img = ( self.img.copy() * background_alpha if background else np.zeros(self.img.shape) ) img = img.astype(np.int32) group_dict = { "both": [True, False], "in": [True], "out": [False], None: [], "None": [], } type_annotation = list(self.nucleus_df["cell_type"]) in_spot = list(self.nucleus_df["in_spot"]) print("Adding cells.") if shape == "cell": d = {t: i for i, t in enumerate(self.type_list)} type_mask = np.zeros( (img.shape[0], img.shape[1], len(self.type_list)), dtype=bool ) for i in tqdm(range(len(self.img_cell))): for j in range(len(self.img_cell[0])): k = self.img_cell[i, j] if ( 0 < k <= len(type_annotation) and type_annotation[k - 1] in d.keys() and in_spot[k - 1] in group_dict[cell] ): type_mask[i, j, d[type_annotation[k - 1]]] = True for i, type_ in tqdm(enumerate(self.type_list)): color = np.array(self.color_dict[type_]).astype(np.int32) img[type_mask[:, :, i]] = color # img.astype(np.int32) for i, type_ in tqdm(enumerate(type_annotation)): if in_spot[i] in group_dict[cell] and type_ in self.color_dict.keys(): color = [int(i) for i in self.color_dict[type_]] if shape == "circle": img = cv.circle( img, self.nucleus_centers[i], circle_size, color, -1 ) elif shape == "nucleus": boundary_temp = self.nucleus_boundary[i].reshape((-1, 1, 2)) img = cv.fillPoly(img, [boundary_temp], color=color) boundary_mask = np.zeros((img.shape[0], img.shape[1]), dtype=bool) for i in tqdm(range(len(type_annotation))): if ( in_spot[i] in group_dict[boundary] and len(self.individual_boundary[i + 1]) > 0 and type_annotation[i] ) in self.color_dict.keys(): index_temp = np.array(self.individual_boundary[i + 1]) boundary_mask[index_temp[:, 0], index_temp[:, 1]] = True img[boundary_mask] = cell_boundary_color if spot: print("Adding spots.") for i in range(len(self.spot_centers)): img = cv.circle( img, self.spot_centers[i], int(self.spot_radius), spot_color, spot_width, ) if save is not None: print("Saving the image.") img1 = img[:, :, [2, 1, 0]] cv.imwrite(save, img1) plt.figure(dpi=dpi) plt.imshow(img)
[docs] def plot_legend( self, save: Optional[str] = None, dpi: int = 300, background: Tuple[int, int, int] = (0, 0, 0), ncol: int = 1, word_color: Tuple[int, int, int] = (255, 255, 255), fontsize: int = 6, horizontal_distance: int = 380, show: bool = False, ) -> None: nrow = len(self.type_list) // ncol if nrow * ncol < len(self.type_list): nrow += 1 img = np.ones((nrow * 60 + 40, ncol * horizontal_distance - 30, 3)) * np.array( background ) img = img.astype(np.int32) fig, ax = plt.subplots(dpi=dpi) plt.imshow(img) i = 0 # col index j = 0 # row index word_color = (word_color[0] / 255, word_color[1] / 255, word_color[2] / 255) for k, v in self.color_dict.items(): ax.add_patch( patches.Rectangle( (i * horizontal_distance + 20, j * 60 + 20), 100, 45, facecolor=np.array(v) / 255, edgecolor="none", ) ) ax.text( i * horizontal_distance + 145, j * 60 + 43, k, va="center", ha="left", fontsize=fontsize, color=word_color, ) j += 1 if j == nrow: j = 0 i += 1 ax.axis("off") if save is not None: plt.savefig(save, bbox_inches="tight", pad_inches=0) if show: plt.show()
[docs] class Plot_Xenium: def __init__( self, Xenium_img: np.ndarray, cell_boundaries: pd.DataFrame, nucleus_boundaries: pd.DataFrame, type_list: List[str], cell_type: List[str], nucleus_centers: np.ndarray, ): """ Args: Xenium_img: Image of Xenium cell_boundaries: Coordinates of the cell boundaries on the image. nucleus_boundaries: Coordinates of the nucleus boundaries on the image. type_list: List of the cell types. cell_type: Annotation. nucleus_centers: Centers of the nuclei. """ self.img = Xenium_img self.type_list = type_list self.cell_type = cell_type self.n_cell = len(cell_type) self.nucleus_centers = nucleus_centers.astype(np.int32) cell_boundaries_temp = cell_boundaries.iloc[:, 1:3].values.astype(np.int32) idx = [[] for _ in range(self.n_cell)] for i in tqdm(range(len(cell_boundaries))): cell_idx = int(cell_boundaries.iloc[i, 0] - 1) idx[cell_idx] += [i] self.cell_boundaries = [ ( cell_boundaries_temp[np.array(idx[i])].reshape((-1, 1, 2)) if len(idx[i]) > 0 else [] ) for i in tqdm(range(self.n_cell)) ] nucleus_boundaries_temp = nucleus_boundaries.iloc[:, 1:3].values.astype( np.int32 ) idx = [[] for _ in range(self.n_cell)] for i in tqdm(range(len(nucleus_boundaries))): cell_idx = int(nucleus_boundaries.iloc[i, 0] - 1) idx[cell_idx] += [i] self.nucleus_boundaries = [ ( nucleus_boundaries_temp[np.array(idx[i])].reshape((-1, 1, 2)) if len(idx[i]) > 0 else [] ) for i in tqdm(range(self.n_cell)) ] self.colors = np.array( list(plt.get_cmap("tab20b").colors) + list(plt.get_cmap("tab20c").colors) ) self.colors = (self.colors * 255).astype(np.int32) self.color_dict = { type_: list(self.colors[i]) for i, type_ in enumerate(type_list) }
[docs] def plot( self, background: bool = False, shape: str = "cell", save: Optional[str] = "Xenium_plot.png", cell_boundaries: bool = False, background_alpha: float = 0.8, circle_size: int = 10, cell_boundary_color: Tuple[int, int, int] = (100, 100, 100), cell_boundary_thickness: int = 2, ) -> None: """ Args: background: If show the background. shape: Which shape to plot? [cell, nucleus, circle] save: If not none, save the figure to the path. cell_boundaries: If plot the cell boundaries. background_alpha: Opacity of the background figure. circle_size: Size of the circle. cell_boundary_color: Color of the cell boundaries. cell_boundary_thickness: Thickness of the cell boundaries. """ img = ( self.img.copy() * background_alpha if background else np.zeros(self.img.shape) ) img = img.astype(np.int32) print("Adding cellls.") for i, type_ in tqdm(enumerate(self.cell_type)): if type_ in self.color_dict.keys(): color = [int(i) for i in self.color_dict[type_]] if shape == "cell": if len(self.cell_boundaries[i]) > 0: img = cv.fillPoly(img, [self.cell_boundaries[i]], color=color) elif shape == "nucleus": if len(self.nucleus_boundaries[i]) > 0: img = cv.fillPoly( img, [self.nucleus_boundaries[i]], color=color ) elif shape == "circle": img = cv.circle( img, self.nucleus_centers[i], circle_size, color, -1 ) if cell_boundaries: for i, type_ in tqdm(enumerate(self.cell_type)): if len(self.cell_boundaries[i]) > 0: img = cv.polylines( img, [self.cell_boundaries[i]], color=cell_boundary_color, thickness=cell_boundary_thickness, isClosed=True, ) if save is not None: print("Saving the image.") img1 = img[:, :, [2, 1, 0]] cv.imwrite(save, img1) plt.imshow(img)