Virtual OpenCell: switch between 2D and 3D virtual staining; 3D shown one confocal slice at a time

app.py

@@ -342,29 +342,38 @@ def build_substitution_scan_section():
 
 def build_virtual_opencell_section():
     """Pick an OpenCell protein and browse CELL-FM's virtual staining of it, laid out like an
-    OpenCell target page: the real cell line's annotations on the left, the images on the right."""
+    OpenCell target page: the real cell line's annotations on the left, the images on the right.
+    "2D" is the 2D projection set, with its map; "3D" the CELL-FM 3D stacks, one z slice at a time."""
 
     def load(gene):
         row = PROTEINS.loc[gene]
         stack, paths = viewer.load_samples(gene, int(row["n_samples"]))
         return row, stack, paths
 
-    def show(gene, index, channel, n_min, n_max, n_gamma, t_min, t_max, t_gamma):
+    def show(gene, index, mode, z, channel, n_min, n_max, n_gamma, t_min, t_max, t_gamma):
+        nucleus, target = (n_min, n_max, n_gamma), (t_min, t_max, t_gamma)
+        if mode == "3D":
+            stack, limits = viewer.load_sample_3d(gene, index)
+            return viewer.render(stack[int(z) - 1], channel, nucleus, target, limits)
         _, stack, _ = load(gene)
-        return viewer.render(stack[index], channel, (n_min, n_max, n_gamma), (t_min, t_max, t_gamma))
+        return viewer.render(stack[index], channel, nucleus, target)
 
-    def sample_file(gene, index):
+    def sample_file(gene, index, mode):
         """The raw sample under a readable name, copied where Gradio is allowed to serve it."""
-        _, _, paths = load(gene)
-        path = os.path.join(tempfile.gettempdir(), f"virtual_opencell_{gene}_{index + 1:04d}.tif")
-        shutil.copyfile(paths[index], path)
+        if mode == "3D":
+            source, name = viewer.sample_file_3d(gene, index), f"virtual_opencell_3d_{gene}_{index + 1:04d}.tif"
+        else:
+            source, name = load(gene)[2][index], f"virtual_opencell_{gene}_{index + 1:04d}.tif"
+        path = os.path.join(tempfile.gettempdir(), name)
+        shutil.copyfile(source, path)
         return path
 
-    def protein_view(gene, *settings):
+    def protein_view(gene, mode, z, *settings):
         """Panel, thumbnails, first sample and its file for one protein."""
-        row, stack, _ = load(gene)
-        image = show(gene, 0, *settings)
-        return viewer.target_html(gene, row), viewer.thumbnails(stack), image, sample_file(gene, 0)
+        row = PROTEINS.loc[gene]
+        thumbs = viewer.thumbnails_3d(gene) if mode == "3D" else viewer.thumbnails(load(gene)[1])
+        image = show(gene, 0, mode, z, *settings)
+        return viewer.target_html(gene, row, mode == "3D"), thumbs, image, sample_file(gene, 0, mode)
 
     def umap_plot(gene):
         try:
@@ -373,21 +382,32 @@ def build_virtual_opencell_section():
             print(f"[virtual opencell] map unavailable: {exc}")
             return None
 
-    def select_protein(gene, channel, n_min, n_max, n_gamma, t_min, t_max, t_gamma):
+    def select_protein(gene, mode, z, *settings):
         try:
-            info, thumbs, image, path = protein_view(gene, channel, n_min, n_max, n_gamma, t_min, t_max, t_gamma)
+            info, thumbs, image, path = protein_view(gene, mode, z, *settings)
         except Exception as exc:
-            raise gr.Error(f"Could not load the images for {gene}: {viewer.load_error(exc)}") from exc
-        return info, gr.Gallery(value=thumbs, selected_index=0), image, 0, path, umap_plot(gene)
-
-    def select_sample(gene, channel, n_min, n_max, n_gamma, t_min, t_max, t_gamma, evt: gr.SelectData):
-        image = show(gene, evt.index, channel, n_min, n_max, n_gamma, t_min, t_max, t_gamma)
-        return image, evt.index, sample_file(gene, evt.index)
-
-    # Shown on first load: the default protein, unless its images cannot be fetched
+            raise gr.Error(f"Could not load the {mode} images for {gene}: {viewer.load_error(exc)}") from exc
+        # there is no 3D map yet: leave the hidden one alone rather than redraw it
+        umap_update = gr.update() if mode == "3D" else umap_plot(gene)
+        return info, gr.Gallery(value=thumbs, selected_index=0), image, 0, path, umap_update
+
+    def select_mode(gene, mode, z, *settings):
+        """2D <-> 3D: the protein's other image set, the slice slider for 3D, the maps for 2D only."""
+        info, samples, image, index, path, _ = select_protein(gene, mode, z, *settings)
+        two_d = mode == "2D"
+        umap_update = gr.update(value=umap_plot(gene), visible=True) if two_d else gr.update(visible=False)
+        return (info, samples, image, index, path, umap_update, gr.update(visible=not two_d),
+                *[gr.update(visible=two_d)] * 3)
+
+    def select_sample(gene, mode, z, channel, n_min, n_max, n_gamma, t_min, t_max, t_gamma, evt: gr.SelectData):
+        image = show(gene, evt.index, mode, z, channel, n_min, n_max, n_gamma, t_min, t_max, t_gamma)
+        return image, evt.index, sample_file(gene, evt.index, mode)
+
+    # Shown on first load: the default protein in 2D, unless its images cannot be fetched
     defaults = ("Both", 0, 100, 1.0, 0, 100, 1.0)
+    slice0 = viewer.DEFAULT_SLICE_3D + 1
     try:
-        info0, thumbs0, image0, file0 = protein_view(viewer.DEFAULT_GENE, *defaults)
+        info0, thumbs0, image0, file0 = protein_view(viewer.DEFAULT_GENE, "2D", slice0, *defaults)
     except Exception as exc:  # e.g. no access to the dataset; the error is shown in the panel
         print(f"[startup] virtual OpenCell preload failed: {exc}")
         info0, thumbs0, image0, file0 = viewer.error_html(viewer.DEFAULT_GENE, exc), None, None, None
@@ -409,11 +429,13 @@ def build_virtual_opencell_section():
         with gr.Row():
             with gr.Column(scale=4):
                 info = gr.HTML(info0)
-                gr.HTML(viewer.UMAP_HTML, padding=False)
+                umap_header = gr.HTML(viewer.UMAP_HTML, padding=False)
                 umap = gr.Plot(value=umap0, show_label=False, elem_classes="vo-umap")
 
             with gr.Column(scale=6):
                 with gr.Row():
+                    mode = gr.Radio(["2D", "3D"], value="2D", label="Virtual staining", scale=0,
+                                    min_width=170, elem_classes="vo-buttons")
                     channel = gr.Radio(viewer.CHANNELS, value=defaults[0], label="Channel",
                                        elem_classes="vo-buttons")
                     download = gr.DownloadButton("Download TIFF", value=file0, size="sm", variant="secondary",
@@ -422,6 +444,9 @@ def build_virtual_opencell_section():
                     value=image0, show_label=False, interactive=False, height=600,
                     show_download_button=True, elem_classes="vo-viewer",
                 )
+                # OpenCell's "2D slice" view: one confocal slice of the stack, picked under the image
+                z = gr.Slider(1, viewer.SLICES_3D, value=slice0, step=1, visible=False,
+                              label=f"Confocal slice (of {viewer.SLICES_3D}, centred on the nucleus)")
                 gr.HTML('<div class="vo-settings-label">Samples</div>', padding=False)
                 samples = gr.Gallery(
                     value=thumbs0, selected_index=0 if thumbs0 else None, show_label=False,
@@ -442,24 +467,26 @@ def build_virtual_opencell_section():
 
         # the full map: every image and every protein named. One fixed figure (~2 MB) sent with the
         # page; the panel map above follows the selection, so this one never has to be re-sent
-        gr.HTML('<div class="vo-target" style="text-align: center; padding-right: 0"><div class="vo-section-header">Full virtual staining map</div>'
-                '<div class="vo-section-caption">All 93,510 generated images, each protein named at the '
-                'centre of its images; zoom in to separate the names</div></div>', padding=False)
-        gr.Plot(value=full_map0, show_label=False, elem_classes="vo-umap-full")
+        full_header = gr.HTML('<div class="vo-target" style="text-align: center; padding-right: 0"><div class="vo-section-header">Full virtual staining map</div>'
+                              '<div class="vo-section-caption">All 93,510 generated images, each protein named at the '
+                              'centre of its images; zoom in to separate the names</div></div>', padding=False)
+        full_map = gr.Plot(value=full_map0, show_label=False, elem_classes="vo-umap-full")
 
         gr.HTML(viewer.FOOTER_HTML)
 
     index = gr.State(0)
     settings = [channel, n_min, n_max, n_gamma, t_min, t_max, t_gamma]
 
-    protein.change(select_protein, inputs=[protein] + settings,
+    protein.change(select_protein, inputs=[protein, mode, z] + settings,
                    outputs=[info, samples, image, index, download, umap])
+    mode.change(select_mode, inputs=[protein, mode, z] + settings,
+                outputs=[info, samples, image, index, download, umap, z, umap_header, full_header, full_map])
 
     # Redraws take milliseconds: no progress overlay, and a dragged slider only renders its latest value
-    samples.select(select_sample, inputs=[protein] + settings, outputs=[image, index, download],
+    samples.select(select_sample, inputs=[protein, mode, z] + settings, outputs=[image, index, download],
                    show_progress="hidden")
-    for control in settings:
-        control.change(show, inputs=[protein, index] + settings, outputs=image,
+    for control in [z] + settings:
+        control.change(show, inputs=[protein, index, mode, z] + settings, outputs=image,
                        show_progress="hidden", trigger_mode="always_last")
 
 

cell_fm/apps/virtual_opencell/viewer.py

@@ -21,6 +21,15 @@ read them from disk instead.
 
 The map of all samples is opencell/vs_umap_embedding.npz in the weights repo: one UMAP point
 per sample (93,510, grouped by protein), embedded with opencell/vit.bin.
+
+The 3D virtual staining (CELL-FM 3D) is in the same repo, 64 samples for every protein:
+
+    OpenCell/virtual_staining_3d/<GENE>/<NNNN>.tif    (48, 2, 192, 192) uint16, z first
+    OpenCell/virtual_staining_3d_thumbnails/<GENE>.jpg
+
+A stack is about 7 MB, so a protein's 64 are fetched one at a time, when opened, and the sample
+strip comes from the thumbnail sheet (notebooks/tools/build_virtual_opencell_3d_thumbnails.py).
+Stacks are shown one z slice at a time, like OpenCell's "2D slice" view. There is no 3D map yet.
 """
 
 import html
@@ -32,10 +41,19 @@ import numpy as np
 import pandas as pd
 import plotly.graph_objects as go
 import tifffile
+from PIL import Image
 
 DATASET_REPO = os.environ.get("VIRTUAL_OPENCELL_REPO", "BoHuangLab/CELL-FM")
 LOCAL_DIR = os.environ.get("VIRTUAL_OPENCELL_LOCAL_DIR", "")
 IMAGE_DIR = "OpenCell/virtual_staining"
+IMAGE_DIR_3D = "OpenCell/virtual_staining_3d"
+THUMBNAIL_DIR_3D = "OpenCell/virtual_staining_3d_thumbnails"
+SAMPLES_3D = 64
+SLICES_3D = 48
+# Crops are centred on the nucleus' 3D centroid, so its centre slice is the default view
+DEFAULT_SLICE_3D = SLICES_3D // 2
+THUMBNAIL_SIZE_3D = 96
+THUMBNAIL_COLUMNS_3D = 8
 
 MODEL_REPO = os.environ.get("CELLFM_MODEL_REPO", "BoHuangLab/CELL-FM")
 UMAP_FILE = "opencell/vs_umap_embedding.npz"
@@ -106,28 +124,61 @@ def load_samples(gene: str, n_samples: int):
     return np.stack([tifffile.imread(p) for p in paths]), tuple(paths)
 
 
-def _scale(channel: np.ndarray, low: float, high: float, gamma: float) -> np.ndarray:
-    """One channel to [0, 1]: auto-contrast, then OpenCell's intensity range (%) and gamma."""
-    lo, hi = np.percentile(channel, AUTO_PERCENTILES)
+def _download(name: str) -> str:
+    if LOCAL_DIR:
+        return os.path.join(LOCAL_DIR, name)
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(repo_id=DATASET_REPO, filename=name, repo_type="dataset")
+
+
+def sample_file_3d(gene: str, index: int) -> str:
+    """Local path of one 3D sample (0-based index), downloading it first unless reading from disk."""
+    return _download(f"{IMAGE_DIR_3D}/{gene}/{index + 1:04d}.tif")
+
+
+@lru_cache(maxsize=32)
+def load_sample_3d(gene: str, index: int):
+    """(stack, limits): stack is (Z, 2, H, W) uint16; limits are each channel's auto-contrast
+    window over the whole stack, so a slice keeps the same contrast as the others."""
+    stack = tifffile.imread(sample_file_3d(gene, index))
+    limits = tuple(tuple(np.percentile(stack[:, c], AUTO_PERCENTILES)) for c in range(stack.shape[1]))
+    return stack, limits
+
+
+@lru_cache(maxsize=64)
+def thumbnails_3d(gene: str) -> list:
+    """The sample strip of a protein's 3D stacks: its centre slices, cut from the thumbnail sheet."""
+    sheet = np.asarray(Image.open(_download(f"{THUMBNAIL_DIR_3D}/{gene}.jpg")).convert("RGB"))
+    size = THUMBNAIL_SIZE_3D
+    return [sheet[r:r + size, c:c + size] for r in range(0, sheet.shape[0], size)
+            for c in range(0, sheet.shape[1], size)]
+
+
+def _scale(channel: np.ndarray, low: float, high: float, gamma: float, limits=None) -> np.ndarray:
+    """One channel to [0, 1]: auto-contrast (over `limits` if given, else this image's own
+    percentiles), then OpenCell's intensity range (%) and gamma."""
+    lo, hi = limits if limits is not None else np.percentile(channel, AUTO_PERCENTILES)
     x = (channel.astype(np.float32) - lo) / max(hi - lo, 1.0)
     low, high = low / 100.0, max(high, low + 1.0) / 100.0
     return np.clip((x - low) / (high - low), 0.0, 1.0) ** gamma
 
 
 def render(sample: np.ndarray, channel: str = "Both",
-           nucleus=(0, 100, 1.0), target=(0, 100, 1.0)) -> np.ndarray:
-    """A (2, H, W) sample as an (H, W, 3) uint8 image; nucleus/target are (min %, max %, gamma)."""
+           nucleus=(0, 100, 1.0), target=(0, 100, 1.0), limits=(None, None)) -> np.ndarray:
+    """A (2, H, W) sample as an (H, W, 3) uint8 image; nucleus/target are (min %, max %, gamma),
+    and limits each channel's auto-contrast window (None: from this image)."""
     rgb = np.zeros(sample.shape[1:] + (3,), dtype=np.float32)
     if channel in ("Nucleus", "Both"):
-        rgb += _scale(sample[0], *nucleus)[..., None] * NUCLEUS_RGB
+        rgb += _scale(sample[0], *nucleus, limits=limits[0])[..., None] * NUCLEUS_RGB
     if channel in ("Target", "Both"):
-        rgb += _scale(sample[1], *target)[..., None] * TARGET_RGB
+        rgb += _scale(sample[1], *target, limits=limits[1])[..., None] * TARGET_RGB
     return (np.clip(rgb, 0.0, 1.0) * 255).astype(np.uint8)
 
 
-def thumbnails(stack: np.ndarray) -> list:
+def thumbnails(stack: np.ndarray, size: int = THUMBNAIL_SIZE) -> list:
     """Both channels at default settings, downsized for the sample strip."""
-    step = max(stack.shape[-1] // THUMBNAIL_SIZE, 1)
+    step = max(stack.shape[-1] // size, 1)
     return [render(sample)[::step, ::step] for sample in stack]
 
 
@@ -227,7 +278,7 @@ def _metadata_item(label: str, value: str) -> str:
             f'<div class="vo-metadata-value">{value}</div></div>')
 
 
-def target_html(gene: str, row: pd.Series) -> str:
+def target_html(gene: str, row: pd.Series, three_d: bool = False) -> str:
     """The left-hand panel of an OpenCell target page: name, identifiers, localization."""
     esc = html.escape
     uniprot, ensg, cid = row["uniprot_id"], row["ensg_id"], row["opencell_cid"]
@@ -258,7 +309,7 @@ def target_html(gene: str, row: pd.Series) -> str:
 <div class="vo-target">
   <div class="vo-target-name">{esc(gene)}</div>
   <div class="vo-protein-description">{esc(name)}</div>
-  <div class="vo-status">CELL-FM virtual staining &middot; {int(row['n_samples'])} samples</div>
+  <div class="vo-status">CELL-FM {"3D " if three_d else ""}virtual staining &middot; {SAMPLES_3D if three_d else int(row['n_samples'])} samples</div>
   <div class="vo-metadata">{''.join(items)}</div>
   <div class="vo-section-header">Protein localization</div>
   <div class="vo-section-caption">OpenCell's annotation of the real cell line; bar height is the grade</div>