Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,162 +1,139 @@
-# app.py
+# streamlit_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,
-)
-
-# ─── FORCE ALL CACHE & CONFIG INTO /tmp ─────────────────────────────────────
+# ─── force all HF/Streamlit caches 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"),
+    ("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)
 
+for d in ("/tmp/streamlit", "/tmp/hf/transformers"):
+    os.makedirs(d, exist_ok=True)
 
-# ─── YOUR HF MODEL REPO ───────────────────────────────────────────────────────
-HF_MODEL_ID = "RakeshNJ12345/Chest-Radiology"
 
+import streamlit as st
+from PIL import Image
+import torch
+import torchvision.transforms as T
+from transformers import ViTModel, T5ForConditionalGeneration, T5Tokenizer
+
+# ─── point at your 1.2 GB model repo, NOT this Space —──────────────────────────
+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)
+    # 1) Vision trunk (fine‐tuned ViT weights under `vit/` in your model repo)
+    vit = ViTModel.from_pretrained(f"{HF_MODEL_ID}/vit").to(device)
 
-    # 2) T5 + tokenizer: same idea, fresh + load YOUR weights
-    t5  = T5ForConditionalGeneration.from_pretrained("t5-base")
+    # 2) T5 + tokenizer (your fine‐tuned report generator)
+    t5  = T5ForConditionalGeneration.from_pretrained(HF_MODEL_ID).to(device)
     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 ─────────────────────────────────────────────────────
+# ─── preprocessing for ViT —────────────────────────────────────────────────────
 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>",
+# ─── Streamlit layout —─────────────────────────────────────────────────────────
+st.set_page_config(page_title="AI Chest X-Ray Report", layout="wide")
+st.markdown("<h1 style='text-align:center;'>AI Chest X-Ray Report</h1>", unsafe_allow_html=True)
+st.markdown("<p style='text-align:center;'>Upload a chest X-ray (PNG/JPG) to generate a report.</p>",
             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 ───────────────────────────────────────────────────────────
+# ─── 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")
+    uploaded = st.file_uploader(
+        "📤 Upload your chest X-ray",
+        type=["png","jpg","jpeg"],
+        label_visibility="collapsed"
+    )
+    if uploaded:
+        st.image(uploaded, width=350,
+                 caption=f"{uploaded.name} — {uploaded.size/1e6:.2f} MB")
         if st.button("▶️ Generate Report"):
-            st.session_state.uploaded = up
+            st.session_state.uploaded = uploaded
             st.session_state.stage    = "report"
             st.experimental_rerun()
 
-
-# ─── REPORT SCREEN ───────────────────────────────────────────────────────────
+# ─── 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]
+        x       = transform(img).unsqueeze(0).to(device)
+        vit_out = 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]
+        # 2) project to T5 hidden size
+        proj       = torch.nn.Linear(vit_out.size(-1), t5.config.d_model).to(device)
+        vision_pf  = proj(vit_out).unsqueeze(1)              # [1,1,d_model]
 
-        # 3) “report:” token embeddings
+        # 3) fixed “report:” prefix
         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 ────────────────────────────────────────────────────────────────
+        txt_emb = t5.encoder.embed_tokens(enc.input_ids)    # [1,L,d_model]
+
+        # 4) build encoder inputs/mask
+        enc_emb  = torch.cat([vision_pf, txt_emb], dim=1)   # [1,1+L,d]
+        enc_mask = torch.cat([
+            torch.ones(1,1,device=device,dtype=torch.long),
+            enc.attention_mask
+        ], dim=1)                                          # [1,1+L]
+
+        # 5) run encoder
+        enc_out = t5.encoder(inputs_embeds=enc_emb,
+                             attention_mask=enc_mask)
+
+        # 6) generate (greedy → no reorder errors)
+        out_ids  = t5.generate(
+            encoder_outputs           = enc_out,
+            encoder_attention_mask    = enc_mask,
+            max_length                = 64,
+            num_beams                 = 1,
+            do_sample                 = False,
+            eos_token_id              = tokenizer.eos_token_id,
+        )
+        diagnosis = tokenizer.decode(out_ids[0], skip_special_tokens=True)
+        confidence = "–"  # you can compute or leave blank
+
+    # ── display side-by-side ─────────────────────────────────────────────────────
     c1, c2 = st.columns(2)
     with c1:
-        st.subheader("Your Uploaded X-ray")
+        st.subheader("Your 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"
-        )
+        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.subheader("AI Report")
         st.markdown(
             f"<div style='background:#e0f7fa;padding:12px;border-radius:6px;'>"
-            f"<strong>Primary Diagnosis</strong><br>{report}</div>",
+            f"<strong>Primary Impression</strong><br>{diagnosis}</div>",
             unsafe_allow_html=True
         )
+        st.markdown(f"**Confidence:** {confidence}")
         if st.button("⬅️ Upload Another"):
             st.session_state.stage = "upload"
             del st.session_state.uploaded
             st.experimental_rerun()
 
+# ─── footer —───────────────────────────────────────────────────────────────────
 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.
+  Powered by a fine-tuned ViT + T5 loaded from your model repo.
 </p>
-""", unsafe_allow_html=True)
+""", unsafe_allow_html=True)