"""EnFormer Deep Ensemble Clustering -- Gradio app (Hugging Face Space ready). Clusters two sets of microscopy images using the EnFormer backbone (Deep Ensemble Clustering, ImageNet-1K pretrained) as a frozen feature extractor, optional interpretable morphology descriptors, and an EnFormer-inspired ensemble consensus clustering (K-Means + Fuzzy C-Means + GMM + Spectral + Agglomerative fused via a co-association matrix). """ import os import sys import tempfile import traceback from typing import List, Optional _HERE = os.path.dirname(os.path.abspath(__file__)) if _HERE not in sys.path: sys.path.insert(0, _HERE) import numpy as np import gradio as gr import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from PIL import Image from data import parse_metadata, ImageItem from pipeline import (cluster_overview_image, assignments_csv, build_features, get_extractor, _ground_truth) from ensemble_cluster import EnsembleClusterer, preprocess_features # Bundled demo images, sampled from the two source folders "1" and "2". EXAMPLES = [ os.path.join(_HERE, "examples", "1"), os.path.join(_HERE, "examples", "2"), ] EXAMPLE_TAGS = ["1", "2"] CLUSTER_COLORS = plt.cm.tab10(np.linspace(0, 1, 10)) # --------------------------------------------------------------------------- viz def _fig_to_image(fig) -> Image.Image: buf = tempfile.NamedTemporaryFile(suffix=".png", delete=False) fig.savefig(buf.name, dpi=110, bbox_inches="tight") plt.close(fig) return Image.open(buf.name) def embedding_plot(emb: np.ndarray, labels: np.ndarray, folders: np.ndarray) -> Image.Image: fig, axes = plt.subplots(1, 2, figsize=(11, 4.4), facecolor="white") for ax, color_by, title in [(axes[0], labels, "By cluster"), (axes[1], folders, "By set (1 / 2)")]: uniq = sorted(set(color_by.tolist())) for i, u in enumerate(uniq): m = color_by == u ax.scatter(emb[m, 0], emb[m, 1], s=22, alpha=0.85, color=CLUSTER_COLORS[i % 10], label=(f"cluster {u}" if title == "By cluster" else f"set {u}")) ax.set_title(title, fontsize=11) ax.set_xticks([]); ax.set_yticks([]) ax.legend(fontsize=8, markerscale=1.2, loc="best", framealpha=0.9) fig.tight_layout() return _fig_to_image(fig) def coassoc_plot(M: np.ndarray, labels: np.ndarray) -> Image.Image: order = np.argsort(labels, kind="stable") fig, ax = plt.subplots(figsize=(5.4, 4.8), facecolor="white") im = ax.imshow(M[np.ix_(order, order)], cmap="magma", vmin=0, vmax=1) ax.set_xticks([]); ax.set_yticks([]) fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="fraction of base clusterings that co-clustered the pair") fig.tight_layout() return _fig_to_image(fig) def metrics_markdown(metrics: dict, n_images: int, n_clusters: int, sizes: list, backbone_used: str) -> str: lines = [ "### Results", f"- Images: {n_images}", f"- Backbone: {backbone_used}", f"- Clusters: {n_clusters}   sizes {sizes}", "", "**Cluster quality**", f"- Silhouette: {metrics.get('silhouette', float('nan')):.3f} (higher better)", f"- Davies-Bouldin: {metrics.get('davies_bouldin', float('nan')):.3f} (lower better)", f"- Calinski-Harabasz: {metrics.get('calinski_harabasz', float('nan')):.0f} (higher better)", ] ext = {k: v for k, v in metrics.items() if k.startswith(("ARI_", "NMI_"))} if ext: lines += ["", "**Alignment with known labels** (ARI / NMI)"] groups = sorted({k.split("_", 1)[1] for k in ext}) for g in groups: ari = metrics.get(f"ARI_{g}", float("nan")) nmi = metrics.get(f"NMI_{g}", float("nan")) lines.append(f"- {g}: {ari:.3f} / {nmi:.3f}") return "\n".join(lines) # --------------------------------------------------------------------------- core def _collect_items(use_examples: bool, files_a, files_b) -> List[ImageItem]: items: List[ImageItem] = [] if use_examples: from data import load_dataset return load_dataset(EXAMPLES, tags=EXAMPLE_TAGS) for tag, files in [("1", files_a), ("2", files_b)]: for f in files or []: path = f.name if hasattr(f, "name") else f items.append(parse_metadata(path, tag)) return items def run_clustering(use_examples, files_a, files_b, backbone, use_morphology, n_clusters, consensus, resolution, progress=gr.Progress()): try: progress(0.02, desc="Collecting images") items = _collect_items(use_examples, files_a, files_b) if len(items) < 6: return (None, None, None, "Please provide at least 6 images across the " "two sets, or use the bundled example dataset.", None) backbone_used = backbone weights = os.environ.get("ENFORMER_WEIGHTS_PATH") try: get_extractor(backbone, weights) except Exception: if backbone == "enformer": backbone_used = "convnext" # graceful fallback if weights unavailable get_extractor("convnext") else: raise def prog(p): progress(0.05 + 0.75 * p, desc="Extracting features") pr = _run_on_items(items, backbone_used, int(resolution), use_morphology, int(n_clusters), consensus, weights, prog) progress(0.85, desc="Rendering") labels = pr.labels sizes = np.bincount(labels).tolist() folders = np.array([it.folder for it in pr.items]) overview = cluster_overview_image(pr.items, labels, per_cluster=8, thumb=130) emb_img = embedding_plot(pr.embedding_2d, labels, folders) \ if pr.embedding_2d is not None else None coassoc = pr.result.co_association heat_img = coassoc_plot(coassoc, labels) if coassoc is not None else None md = metrics_markdown(pr.metrics, len(pr.items), len(sizes), sizes, "EnFormer-Small (ImageNet-1K)" if backbone_used == "enformer" else "ConvNeXt-Tiny (fallback)") csv = assignments_csv(pr.items, labels) csv_path = os.path.join(tempfile.mkdtemp(), "cluster_assignments.csv") with open(csv_path, "w") as fh: fh.write(csv) return overview, emb_img, heat_img, md, csv_path except Exception as e: tb = traceback.format_exc() return None, None, None, f"**Error:** {e}\n\n```\n{tb[-1500:]}\n```", None def _run_on_items(items, backbone, resolution, use_morphology, n_clusters, consensus, weights, prog): """Run the pipeline directly on a list of items (uploads or examples).""" from pipeline import PipelineResult from sklearn.cluster import SpectralClustering, AgglomerativeClustering feat, _ = build_features(items, backbone, resolution, 1, use_morphology, weights, progress=prog) Z = preprocess_features(feat, pca_dim=min(12, feat.shape[1] - 1, len(items) - 1), whiten=False, l2=True) gt = _ground_truth(items) ec = EnsembleClusterer(n_clusters=n_clusters, base_methods=("kmeans", "fuzzy", "gmm", "spectral", "agglomerative"), n_runs=10, subspace_frac=0.85, k_jitter=0) def avg_link(M, k): d = 1.0 - M; np.fill_diagonal(d, 0.0) return AgglomerativeClustering(k, metric="precomputed", linkage="average").fit_predict(d) def spectral(M, k): return SpectralClustering(k, affinity="precomputed", assign_labels="kmeans", random_state=0).fit_predict(M) ec._consensus = spectral if consensus == "spectral" else avg_link res = ec.fit(Z, ground_truth=gt, compute_2d=True) from pipeline import _merge_tiny_clusters labels = _merge_tiny_clusters(Z, res.labels, min_size=max(2, int(0.02 * len(items)))) if not np.array_equal(labels, res.labels): res.labels = labels res.metrics = EnsembleClusterer.compute_metrics(Z, labels, gt) return PipelineResult(items=items, labels=res.labels, result=res, metrics=res.metrics, embedding_2d=res.embedding_2d) # --------------------------------------------------------------------------- UI DESCRIPTION = """ # EnFormer Deep Ensemble Clustering Unsupervised clustering of two image sets. Images are encoded with the EnFormer backbone (ImageNet-pretrained) and grouped by an ensemble of base clusterers (K-Means, Fuzzy C-Means, Gaussian-Mixture, Spectral, Agglomerative) fused through a co-association matrix. """ THEME = gr.themes.Base( primary_hue=gr.themes.colors.gray, secondary_hue=gr.themes.colors.gray, neutral_hue=gr.themes.colors.gray, font=["system-ui", "Arial", "sans-serif"], ).set( body_background_fill="#ffffff", body_text_color="#000000", body_text_color_subdued="#444444", background_fill_primary="#ffffff", background_fill_secondary="#ffffff", block_background_fill="#ffffff", block_border_color="#d0d0d0", block_label_text_color="#000000", block_title_text_color="#000000", border_color_primary="#d0d0d0", input_background_fill="#ffffff", input_border_color="#c0c0c0", button_primary_background_fill="#000000", button_primary_background_fill_hover="#333333", button_primary_text_color="#ffffff", button_secondary_background_fill="#ffffff", button_secondary_text_color="#000000", button_secondary_border_color="#000000", ) # Force a white background / black text even when the viewer's system (or the # Hugging Face host page) is in dark mode. Gradio's dark theme sets many CSS # variables via the `.dark` class, so we override the full relevant set back to # light values -- no page reload needed. _LIGHT_VARS = """ color-scheme: light; --body-background-fill:#ffffff; --body-text-color:#000000; --body-text-color-subdued:#4b4b4b; --background-fill-primary:#ffffff; --background-fill-secondary:#f6f6f6; --border-color-primary:#d4d4d4; --border-color-accent:#bcbcbc; --block-background-fill:#ffffff; --block-border-color:#d4d4d4; --block-label-background-fill:#ffffff; --block-label-text-color:#000000; --block-title-text-color:#000000; --block-info-text-color:#4b4b4b; --panel-background-fill:#ffffff; --panel-border-color:#d4d4d4; --input-background-fill:#ffffff; --input-background-fill-focus:#ffffff; --input-border-color:#c4c4c4; --input-border-color-focus:#000000; --input-placeholder-color:#8a8a8a; --button-primary-background-fill:#000000; --button-primary-background-fill-hover:#333333; --button-primary-text-color:#ffffff; --button-primary-border-color:#000000; --button-secondary-background-fill:#f2f2f2; --button-secondary-background-fill-hover:#e6e6e6; --button-secondary-text-color:#000000; --button-secondary-border-color:#c4c4c4; --checkbox-background-color:#ffffff; --checkbox-background-color-hover:#f0f0f0; --checkbox-background-color-selected:#000000; --checkbox-border-color:#999999; --checkbox-border-color-hover:#666666; --checkbox-border-color-focus:#000000; --checkbox-border-color-selected:#000000; --checkbox-label-background-fill:#f2f2f2; --checkbox-label-background-fill-hover:#e9e9e9; --checkbox-label-background-fill-selected:#e9e9e9; --checkbox-label-text-color:#000000; --checkbox-label-text-color-selected:#000000; --checkbox-label-border-color:#c4c4c4; --slider-color:#000000; --table-even-background-fill:#ffffff; --table-odd-background-fill:#f6f6f6; --table-border-color:#d4d4d4; --accordion-text-color:#000000; """ CSS = """ .gradio-container { max-width:1180px !important; } footer { display:none !important; } :root, .dark, gradio-app, .gradio-container { """ + _LIGHT_VARS + """ } body, .gradio-container, .dark { background:#ffffff !important; color:#000000 !important; } /* selected radio/checkbox option labels keep black text (dark-mode used near-white) */ label.selected, label.selected *, .selected > span { color:#000000 !important; } """ def build_demo(): with gr.Blocks(title="EnFormer Ensemble Clustering", theme=THEME, css=CSS) as demo: gr.Markdown(DESCRIPTION) with gr.Row(): with gr.Column(scale=1): use_examples = gr.Checkbox(value=True, label="Use bundled example dataset") with gr.Accordion("Upload your own two image sets", open=False): files_a = gr.File(label="Set 1", file_count="multiple", file_types=["image"]) files_b = gr.File(label="Set 2", file_count="multiple", file_types=["image"]) backbone = gr.Radio(["enformer", "convnext"], value="enformer", label="Feature backbone") use_morphology = gr.Checkbox(value=True, label="Add morphology features") n_clusters = gr.Slider(2, 10, value=5, step=1, label="Number of clusters (K)") consensus = gr.Radio(["spectral", "average_linkage"], value="spectral", label="Consensus function") resolution = gr.Radio([224, 320], value=224, label="Input resolution") run_btn = gr.Button("Run clustering", variant="primary") with gr.Column(scale=2): metrics_md = gr.Markdown() csv_out = gr.File(label="Cluster assignments (CSV)") clusters_out = gr.Image(label="Clusters (representative images per cluster)", type="pil", show_label=True) with gr.Row(): emb_out = gr.Image(label="2-D feature embedding", type="pil") heat_out = gr.Image(label="Consensus co-association matrix", type="pil") run_btn.click( run_clustering, inputs=[use_examples, files_a, files_b, backbone, use_morphology, n_clusters, consensus, resolution], outputs=[clusters_out, emb_out, heat_out, metrics_md, csv_out], ) return demo if __name__ == "__main__": demo = build_demo() demo.queue().launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))