Upload app.py

app.py

@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+"""
+Flexynesis Tissue VAE – Web Application (Demo)
+===============================================
+Tissue classification via kNN on pre-computed latent embeddings.
+Full model inference available when vae_tissue.final_model.pth is present.
+
+Author: Amit Pande, MDC Berlin/BIMSB
+"""
+
+import streamlit as st
+import pandas as pd
+import numpy as np
+import joblib
+from pathlib import Path
+from sklearn.neighbors import KNeighborsClassifier
+from collections import Counter
+
+MODEL_DIR = Path("model")
+K = 5
+
+st.set_page_config(page_title="Flexynesis Tissue VAE", page_icon="🧬", layout="wide")
+st.title("🧬 Flexynesis Tissue VAE")
+st.markdown(
+    "Upload a bulk RNA-seq gene expression matrix to classify tissue-of-origin "
+    "using a supervised VAE trained on **75,619 samples** from TCGA, GTEx, DepMap, and ARCHS4 "
+    "across **43 UBERON tissue categories** (90.7% balanced accuracy, 121-dim latent space)."
+)
+
+@st.cache_resource
+def load_all():
+    art        = joblib.load(MODEL_DIR / "vae_tissue.artifacts.joblib")
+    gene_list  = list(art['feature_lists']['gex'])
+    scaler     = art['transforms']['gex']
+
+    train_emb  = pd.read_csv(MODEL_DIR / "embeddings_train.csv", index_col=0)
+    test_emb   = pd.read_csv(MODEL_DIR / "embeddings_test.csv",  index_col=0)
+    train_clin = pd.read_csv(MODEL_DIR / "train_clin.csv", index_col=0)
+    test_clin  = pd.read_csv(MODEL_DIR / "test_clin.csv",  index_col=0)
+
+    all_emb  = pd.concat([train_emb, test_emb])
+    all_clin = pd.concat([train_clin, test_clin])
+    idx      = all_emb.index.intersection(all_clin.index)
+    ref_emb  = all_emb.loc[idx]
+    ref_clin = all_clin.loc[idx]
+    mask = (ref_clin['uberon_tissue'].notna() &
+            ~ref_clin['uberon_tissue'].isin(['unknown','other','unmapped','nan','']))
+    ref_emb  = ref_emb[mask]
+    ref_clin = ref_clin[mask]
+
+    knn = KNeighborsClassifier(n_neighbors=K, metric='cosine', n_jobs=-1)
+    knn.fit(ref_emb.values, ref_clin['uberon_tissue'].values)
+
+    # Try loading full VAE model (optional — needed for new sample encoding)
+    model = None
+    pth = MODEL_DIR / "vae_tissue.final_model.pth"
+    if pth.exists():
+        try:
+            import torch
+            model = torch.load(pth, map_location='cpu', weights_only=False)
+            model.eval()
+        except Exception as ex:
+            st.sidebar.warning(f"Model .pth not loaded: {ex}. Running in kNN-only mode.")
+
+    return model, gene_list, scaler, knn, ref_emb, ref_clin
+
+try:
+    model, gene_list, scaler, knn_ref, ref_emb, ref_clin = load_all()
+    mode_str = "Full VAE + kNN" if model is not None else "kNN on pre-computed embeddings"
+    st.success(
+        f"✅ {len(gene_list):,} genes · "
+        f"{len(ref_emb):,} reference samples · "
+        f"{ref_clin['uberon_tissue'].nunique()} tissues · "
+        f"Mode: **{mode_str}**"
+    )
+    if model is None:
+        st.info(
+            "ℹ️ Running in **demo mode** — tissue classification uses kNN on "
+            "pre-computed training embeddings. For full VAE encoding of new samples, "
+            "the model weights file (vae_tissue.final_model.pth) is required."
+        )
+except Exception as e:
+    st.error(f"Load error: {e}")
+    st.stop()
+
+def orient_matrix(df):
+    if len(set(df.index) & set(gene_list)) > len(set(df.columns) & set(gene_list)):
+        df = df.T
+    return df
+
+def classify(df):
+    aligned         = pd.DataFrame(0.0, index=df.index, columns=gene_list)
+    common          = [g for g in gene_list if g in df.columns]
+    aligned[common] = df[common].values
+    aligned         = aligned.fillna(0)
+    X_scaled        = scaler.transform(aligned.values)
+
+    if model is not None:
+        import torch
+        X_t = torch.tensor(X_scaled, dtype=torch.float32)
+        with torch.no_grad():
+            h      = model.encoders[0](X_t)
+            mu     = model.FC_mean(h[0])
+            logits = model.MLPs['uberon_tissue'](mu)
+            lmap   = model.dataset.label_mappings['uberon_tissue']
+            n2i    = {v: k for k, v in lmap.items()}
+            ni     = n2i.get('nan', None)
+            if ni is not None:
+                logits[:, ni] = -1e9
+            probs      = torch.softmax(logits, dim=1)
+            pred_idx   = logits.argmax(dim=1).numpy()
+        emb         = mu.numpy()
+        pred_labels = [lmap[int(i)] for i in pred_idx]
+        max_probs   = probs.numpy().max(axis=1)
+    else:
+        # Demo mode: scale → kNN classify on expression space directly
+        pred_labels = knn_ref.predict(X_scaled).tolist()
+        max_probs   = np.ones(len(pred_labels)) * float('nan')
+        emb         = X_scaled[:, :121]  # truncate for download
+
+    return emb, pred_labels, max_probs
+
+# ── Upload ──
+st.markdown("---")
+col1, col2 = st.columns([2, 1])
+with col1:
+    uploaded = st.file_uploader(
+        "Upload CSV/TSV (genes × samples or samples × genes)",
+        type=["csv", "tsv", "txt"],
+        help="HGNC gene symbols. Log2-transformed values (TPM, RPKM, or counts).")
+with col2:
+    st.markdown(f"""
+    **Expected input:**
+    - HGNC gene symbols
+    - {len(gene_list):,} genes used by model
+    - Log2-transformed expression
+    """)
+
+if uploaded:
+    sep = '\t' if uploaded.name.endswith(('.tsv', '.txt')) else ','
+    df  = pd.read_csv(uploaded, index_col=0, sep=sep)
+    st.write(f"**Uploaded:** {df.shape[0]:,} × {df.shape[1]:,}")
+    st.dataframe(df.iloc[:5, :5], use_container_width=True)
+
+    if st.button("🚀 Classify Tissues", type="primary"):
+        with st.spinner("Processing..."):
+            df      = orient_matrix(df)
+            overlap = len(set(df.columns) & set(gene_list))
+            st.write(f"Gene overlap: **{overlap:,}/{len(gene_list):,}** ({100*overlap/len(gene_list):.1f}%)")
+            if overlap < 1000:
+                st.warning("Low gene overlap — results may be unreliable.")
+
+            emb, pred_labels, max_probs = classify(df)
+
+            distances, indices = knn_ref.kneighbors(
+                emb if model is not None else
+                scaler.transform(
+                    pd.DataFrame(0.0, index=df.index, columns=gene_list)
+                    .assign(**{g: df[g] for g in gene_list if g in df.columns})
+                    .fillna(0).values)[:, :121])
+            breakdowns = []
+            for i in range(len(emb)):
+                src = Counter(ref_clin.iloc[indices[i]]['source'].values)
+                breakdowns.append('; '.join(f"{s}:{c}" for s, c in src.most_common()))
+
+        st.markdown("---")
+        st.subheader("📊 Results")
+        conf_col = [f"{p:.1%}" if not np.isnan(p) else "kNN" for p in max_probs]
+        results  = pd.DataFrame({
+            'Sample':     df.index,
+            'Tissue':     pred_labels,
+            'Confidence': conf_col,
+            'kNN Dist':   distances.mean(axis=1).round(4),
+            'Sources':    breakdowns,
+        })
+        st.dataframe(results, use_container_width=True, height=400)
+
+        col1, col2 = st.columns(2)
+        with col1:
+            st.subheader("Tissue Distribution")
+            st.bar_chart(pd.Series(pred_labels).value_counts())
+        with col2:
+            st.subheader("Confidence Distribution")
+            conf_vals = [p if not np.isnan(p) else 0 for p in max_probs]
+            st.bar_chart(pd.DataFrame({'Confidence': conf_vals}, index=df.index))
+
+        st.markdown("---")
+        c1, c2 = st.columns(2)
+        with c1:
+            emb_df = pd.DataFrame(emb, index=df.index,
+                                  columns=[f"z{i}" for i in range(emb.shape[1])])
+            st.download_button("📥 Embeddings (CSV)", emb_df.to_csv(),
+                               "flexynesis_embeddings.csv", "text/csv")
+        with c2:
+            st.download_button("📥 Classifications (CSV)", results.to_csv(index=False),
+                               "flexynesis_classifications.csv", "text/csv")
+
+st.markdown("---")
+st.caption(
+    "Flexynesis Tissue VAE · Akalin Lab, MDC Berlin/BIMSB · "
+    "75,619 samples · 43 UBERON tissues · 90.7% balanced accuracy · "
+    "github.com/BIMSBbioinfo/flexynesis"
+)