Merge branch 'main' of github.com:xing-lab-pitt/LiveCellTracker-dev

livecellx/core/pl_utils.py

@@ -1,7 +1,91 @@
 from mpl_toolkits.axes_grid1 import make_axes_locatable
+import numpy as np
+from typing import List
+from livecellx.core.single_cell import SingleCellStatic
+from livecellx.core.utils import crop_or_pad_img
+from livecellx.livecell_logger import main_info, main_warning
 
 
 def add_colorbar(im, ax, fig):
     divider = make_axes_locatable(ax)
     cax = divider.append_axes("right", size="3%", pad=0.05)
     fig.colorbar(im, cax=cax, orientation="vertical")
+
+
+def viz_embedding_region(
+    embedding,
+    scs: List[SingleCellStatic],
+    padding=40,
+    x_range=(float("-inf"), float("inf")),
+    y_range=(float("-inf"), float("inf")),
+    title="Single cell in Embedding Space",
+    max_crops=10,
+    fix_dims=None,
+    randomly_select=True,
+    sort_by_x=True,
+    dpi=300,
+    show_mask=False,
+):
+    _crops = []
+    mask_crops = []
+    _crop_coords = []
+    if fix_dims is not None:
+        main_info(f"Setting: fix_dims={fix_dims}")
+
+    indices = range(len(scs))
+    if randomly_select:
+        indices = np.random.choice(len(scs), len(scs), replace=False)
+
+    for idx in indices:
+        _sc = scs[idx]
+        # print("embedding[idx]", embedding[idx])
+        # Embedding in range
+        if x_range[0] <= embedding[idx][0] <= x_range[1] and y_range[0] <= embedding[idx][1] <= y_range[1]:
+            # _sc.show(crop=True, padding=20)
+            # plt.show()
+            if show_mask:
+                mask_crop = _sc.get_mask_crop(padding=padding)
+            img_crop = _sc.get_img_crop(padding=padding)
+
+            if not (fix_dims is None):
+                img_crop = crop_or_pad_img(img_crop, fix_dims=fix_dims)
+                if show_mask:
+                    mask_crop = crop_or_pad_img(mask_crop, fix_dims=fix_dims)
+            _crops.append(img_crop)
+            if show_mask:
+                mask_crops.append(mask_crop)
+            _crop_coords.append(embedding[idx])
+        if len(_crops) >= max_crops:
+            break
+
+    if len(_crops) == 0:
+        main_info(f"No single cell found in the embedding space")
+        return
+
+    if sort_by_x:
+        main_info(f"Sorting crops by x")
+        sort_indices = np.argsort(np.array(_crop_coords)[:, 0])
+        _crops = [_crops[i] for i in sort_indices]
+        if show_mask:
+            mask_crops = [mask_crops[i] for i in sort_indices]
+        _crop_coords = [_crop_coords[i] for i in sort_indices]
+
+    # Visualize crops on one orw
+    import matplotlib.pyplot as plt
+
+    if show_mask:
+        fig, axes = plt.subplots(2, len(_crops), figsize=(len(_crops) * 4, 10), dpi=dpi)
+    else:
+        fig, axes = plt.subplots(1, len(_crops), figsize=(len(_crops) * 4, 5), dpi=dpi)
+        axes = [axes]
+
+    for idx in range(len(_crops)):
+        axes[0][idx].imshow(_crops[idx])
+        axes[0][idx].axis("off")
+        if show_mask:
+            axes[1][idx].imshow(mask_crops[idx])
+            axes[1][idx].axis("off")
+
+    fig.suptitle(title, fontsize=20)
+    fig.tight_layout(rect=[0, 0.03, 1, 0.95], w_pad=0, h_pad=0)
+    return fig, axes