Update streamlit_app.py

streamlit_app.py

@@ -1,40 +1,162 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+# app.py
+import os
+import torch
+import torch.nn as nn
+import torchvision.transforms as T
 import streamlit as st
+from PIL import Image
+from transformers import (
+    ViTConfig, ViTModel,
+    T5ForConditionalGeneration,
+    T5Tokenizer,
+)
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+# ─── FORCE ALL CACHE & CONFIG INTO /tmp ─────────────────────────────────────
+for ENV, VAL in [
+    ("HOME",               "/tmp"),
+    ("XDG_CONFIG_HOME",    "/tmp"),
+    ("STREAMLIT_HOME",     "/tmp"),
+    ("XDG_CACHE_HOME",     "/tmp"),
+    ("HF_HOME",            "/tmp/hf"),
+    ("TRANSFORMERS_CACHE", "/tmp/hf/transformers"),
+]:
+    os.environ[ENV] = VAL
+os.makedirs("/tmp/streamlit", exist_ok=True)
+os.makedirs("/tmp/hf/transformers", exist_ok=True)
+
+
+# ─── YOUR HF MODEL REPO ───────────────────────────────────────────────────────
+HF_MODEL_ID = "RakeshNJ12345/Chest-Radiology"
+
+
+@st.cache_resource(show_spinner=False)
+def load_models():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # 1) VIT: load its config, build fresh, then we'll load YOUR weights into it
+    vit_cfg = ViTConfig.from_pretrained("google/vit-base-patch16-224")
+    vit     = ViTModel(vit_cfg)
+
+    # 2) T5 + tokenizer: same idea, fresh + load YOUR weights
+    t5  = T5ForConditionalGeneration.from_pretrained("t5-base")
+    tok = T5Tokenizer.from_pretrained(HF_MODEL_ID)
+
+    # 3) grab the single combined file from your repo
+    state = torch.hub.load_state_dict_from_url(
+        f"https://huggingface.co/{HF_MODEL_ID}/resolve/main/pytorch_model.bin",
+        map_location="cpu", check_hash=True
+    )
+
+    # 4) split into vit vs t5 state_dicts
+    vit_state = {k[len("vit."):]: v for k,v in state.items() if k.startswith("vit.")}
+    t5_state  = {k[len("t5."):]:  v for k,v in state.items() if k.startswith("t5.")}
+
+    # 5) load them
+    vit.load_state_dict(vit_state, strict=False)
+    t5.load_state_dict(t5_state, strict=False)
+
+    # 6) move to device & eval
+    vit.to(device).eval()
+    t5.to(device).eval()
+
+    return device, vit, t5, tok
+
+
+device, vit, t5, tokenizer = load_models()
+
+
+# ─── IMAGE PREPROCESSING ─────────────────────────────────────────────────────
+transform = T.Compose([
+    T.Resize((224, 224)),
+    T.ToTensor(),
+    T.Normalize(mean=0.5, std=0.5),
+])
+
+
+# ─── STREAMLIT LAYOUT ────────────────────────────────────────────────────────
+st.set_page_config(page_title="Radiology Report Analysis", layout="wide")
+st.markdown("<h1 style='text-align:center;'>🩺 Radiology Report Analysis</h1>",
+            unsafe_allow_html=True)
+st.markdown(
+    "<p style='text-align:center;'>Upload a chest X-ray (PNG/JPG) to generate an AI report.</p>",
+    unsafe_allow_html=True
+)
+
+if "stage" not in st.session_state:
+    st.session_state.stage = "upload"
+
+
+# ─── UPLOAD SCREEN ───────────────────────────────────────────────────────────
+if st.session_state.stage == "upload":
+    up = st.file_uploader("", type=["png","jpg","jpeg"], label_visibility="collapsed")
+    if up:
+        st.image(up, width=350, caption=f"{up.name} — {up.size/1e6:.2f} MB")
+        if st.button("▶️ Generate Report"):
+            st.session_state.uploaded = up
+            st.session_state.stage    = "report"
+            st.experimental_rerun()
+
+
+# ─── REPORT SCREEN ───────────────────���───────────────────────────────────────
+elif st.session_state.stage == "report":
+    img = Image.open(st.session_state.uploaded).convert("RGB")
+
+    with st.spinner("🔎 Analyzing…"):
+        # 1) ViT features
+        x      = transform(img).unsqueeze(0).to(device)
+        vfeat  = vit(pixel_values=x).pooler_output         # [1,768]
+
+        # 2) project into T5’s hidden size
+        proj   = nn.Linear(vfeat.size(-1), t5.config.d_model).to(device)
+        prefix = proj(vfeat).unsqueeze(1)                  # [1,1,d_model]
+
+        # 3) “report:” token embeddings
+        enc     = tokenizer("report:", return_tensors="pt").to(device)
+        txt_emb = t5.encoder.embed_tokens(enc.input_ids)   # [1,L,d_model]
+
+        # 4) concat + mask
+        emb     = torch.cat([prefix, txt_emb], dim=1)      # [1,1+L,d]
+        mask    = torch.cat([
+                     torch.ones(1,1,device=device), 
+                     enc.attention_mask
+                  ], dim=1)                               # [1,1+L]
+
+        # 5) encode + generate
+        enc_out = t5.encoder(inputs_embeds=emb, attention_mask=mask)
+        ids     = t5.generate(
+                     encoder_outputs            = enc_out,
+                     encoder_attention_mask     = mask,
+                     max_length                 = 64,
+                     num_beams                  = 1,
+                     do_sample                  = False,
+                     eos_token_id               = tokenizer.eos_token_id,
+                 )
+        report  = tokenizer.decode(ids[0], skip_special_tokens=True)
+
+    # ── DISPLAY ────────────────────────────────────────────────────────────────
+    c1, c2 = st.columns(2)
+    with c1:
+        st.subheader("Your Uploaded X-ray")
+        st.image(img, use_column_width=True)
+        st.markdown(
+            f"**File:** {st.session_state.uploaded.name}  \n"
+            f"**Size:** {st.session_state.uploaded.size/1e6:.2f} MB"
+        )
+    with c2:
+        st.subheader("AI Diagnosis & Report")
+        st.markdown(
+            f"<div style='background:#e0f7fa;padding:12px;border-radius:6px;'>"
+            f"<strong>Primary Diagnosis</strong><br>{report}</div>",
+            unsafe_allow_html=True
+        )
+        if st.button("⬅️ Upload Another"):
+            st.session_state.stage = "upload"
+            del st.session_state.uploaded
+            st.experimental_rerun()
+
+st.markdown("""
+<hr style='margin:2em 0;'>
+<p style='font-size:0.8em;color:gray;text-align:center;'>
+  Powered by your fine-tuned ViT➔T5, both coming from a single pytorch_model.bin in Chest-Radiology.
+</p>
+""", unsafe_allow_html=True)