Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,124 +1,180 @@
-# app.py or streamlit_app.py
+# streamlit_app.py
 
+# ──── SET ENVIRONMENT VARIABLES BEFORE ANY IMPORTS ──────────────────────────────
 import os
-import streamlit as st
-from PIL import Image
-import torch
-import torchvision.transforms as T
-import pydicom
-import numpy as np
-from transformers import ViTFeatureExtractor, AutoTokenizer, VisionEncoderDecoderModel
-
-# ─── FORCE ALL CACHE & CONFIG INTO /tmp ────────────────────────────────────────
-# must come before streamlit or transformers imports write any files
-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"),
+import tempfile
+
+# Create a dedicated cache directory
+CACHE_DIR = "/tmp/hf_cache"
+os.makedirs(CACHE_DIR, exist_ok=True)
+
+# Set all relevant environment variables
+os.environ.update({
+    "HOME": "/tmp",
+    "XDG_CONFIG_HOME": "/tmp",
+    "STREAMLIT_HOME": "/tmp/streamlit",
+    "XDG_CACHE_HOME": CACHE_DIR,
+    "HF_HOME": f"{CACHE_DIR}/huggingface",
+    "TRANSFORMERS_CACHE": f"{CACHE_DIR}/transformers",
+    "HF_HUB_CACHE": f"{CACHE_DIR}/huggingface_hub",
+    "HUGGINGFACE_HUB_CACHE": f"{CACHE_DIR}/huggingface_hub"
+})
+
+# Create all cache directories explicitly
+for path in [
+    "/tmp/streamlit",
+    f"{CACHE_DIR}/huggingface",
+    f"{CACHE_DIR}/transformers",
+    f"{CACHE_DIR}/huggingface_hub"
 ]:
-    os.environ[ENV] = VAL
+    os.makedirs(path, exist_ok=True)
 
-os.makedirs("/tmp/streamlit", exist_ok=True)
-os.makedirs("/tmp/hf/transformers", exist_ok=True)
+# ──── NOW IMPORT OTHER LIBRARIES ───────────────────────────────────────────────
+import json
+import torch
+import torch.nn as nn
+import torchvision.transforms as T
+import streamlit as st
+from PIL import Image
+from transformers import ViTModel, T5ForConditionalGeneration, T5Tokenizer
+from huggingface_hub import hf_hub_download
 
-# ─── YOUR MODEL ID ─────────────────────────────────────────────────────────────
+# ──── MODEL DEFINITION ─────────────────────────────────────────────────────────
 MODEL_ID = "RakeshNJ12345/Chest-Radiology"
 
+class TwoViewVisionReportModel(nn.Module):
+    def __init__(self, vit: ViTModel, t5: T5ForConditionalGeneration, tokenizer: T5Tokenizer):
+        super().__init__()
+        self.vit = vit
+        self.proj_f = nn.Linear(vit.config.hidden_size, t5.config.d_model)
+        self.proj_l = nn.Linear(vit.config.hidden_size, t5.config.d_model)
+        self.tokenizer = tokenizer
+        self.t5 = t5
+
+    def generate(self, img: torch.Tensor, max_length: int = 64) -> torch.Tensor:
+        device = img.device
+        vf = self.vit(pixel_values=img).pooler_output
+        pf = self.proj_f(vf).unsqueeze(1)
+        prefix = pf  # single-view only
+
+        enc = self.tokenizer("report:", return_tensors="pt").to(device)
+        txt_emb = self.t5.encoder.embed_tokens(enc.input_ids)
+
+        enc_emb = torch.cat([prefix, txt_emb], dim=1)
+        enc_mask = torch.cat([
+            torch.ones(1, 1, device=device, dtype=torch.long),
+            enc.attention_mask
+        ], dim=1)
+
+        enc_out = self.t5.encoder(
+            inputs_embeds=enc_emb,
+            attention_mask=enc_mask
+        )
+
+        out_ids = self.t5.generate(
+            encoder_outputs=enc_out,
+            encoder_attention_mask=enc_mask,
+            max_length=max_length,
+            num_beams=1,
+            do_sample=False,
+            eos_token_id=self.tokenizer.eos_token_id,
+        )
+        return out_ids
+
+# ──── MODEL LOADING WITH ROBUST CACHE HANDLING ─────────────────────────────────
 @st.cache_resource(show_spinner=False)
-def load_model():
+def load_models():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    fe = ViTFeatureExtractor.from_pretrained(MODEL_ID)
-    tok = AutoTokenizer.from_pretrained(MODEL_ID)
-    pipe = VisionEncoderDecoderModel.from_pretrained(MODEL_ID).to(device)
-    return device, fe, tok, pipe
+    
+    # Ensure cache directories exist
+    for path in [
+        f"{CACHE_DIR}/huggingface",
+        f"{CACHE_DIR}/transformers",
+        f"{CACHE_DIR}/huggingface_hub"
+    ]:
+        os.makedirs(path, exist_ok=True)
 
-device, feat_ext, tokenizer, model = load_model()
+    # Download config with explicit cache
+    cfg_path = hf_hub_download(
+        repo_id=MODEL_ID,
+        filename="config.json",
+        repo_type="model",
+        cache_dir=f"{CACHE_DIR}/huggingface_hub",
+        local_files_only=False,
+        force_download=True
+    )
+    cfg = json.load(open(cfg_path, "r"))
+
+    # Load models with explicit cache directories
+    vit = ViTModel.from_pretrained(
+        "google/vit-base-patch16-224",
+        ignore_mismatched_sizes=True,
+        cache_dir=f"{CACHE_DIR}/transformers"
+    ).to(device)
+    
+    t5 = T5ForConditionalGeneration.from_pretrained(
+        "t5-base",
+        cache_dir=f"{CACHE_DIR}/transformers"
+    ).to(device)
+    
+    tok = T5Tokenizer.from_pretrained(
+        MODEL_ID,
+        cache_dir=f"{CACHE_DIR}/transformers"
+    )
 
-# ─── IMAGE PREPROCESSING ───────────────────────────────────────────────────────
+    # Load combined model
+    model = TwoViewVisionReportModel(vit, t5, tok).to(device)
+    
+    ckpt_path = hf_hub_download(
+        repo_id=MODEL_ID,
+        filename="pytorch_model.bin",
+        repo_type="model",
+        cache_dir=f"{CACHE_DIR}/huggingface_hub",
+        local_files_only=False,
+        force_download=True
+    )
+    
+    state = torch.load(ckpt_path, map_location=device)
+    model.load_state_dict(state)
+    return device, model, tok
+
+# ──── APP INTERFACE ───────────────────────────────────────────────────────────
+device, model, tokenizer = load_models()
 transform = T.Compose([
     T.Resize((224, 224)),
     T.ToTensor(),
     T.Normalize(mean=0.5, std=0.5),
 ])
 
-def load_image(uploaded_file):
-    """Handle .dcm or normal images uniformly, returns a PIL RGB image."""
-    name = uploaded_file.name.lower()
-    if name.endswith(".dcm"):
-        ds = pydicom.dcmread(uploaded_file)
-        arr = ds.pixel_array.astype(np.float32)
-        # normalize to 0–255
-        arr = (arr - arr.min()) / (arr.max() - arr.min()) * 255.0
-        arr = arr.astype(np.uint8)
-        # if monochrome, convert to RGB by stacking
-        if arr.ndim == 2:
-            arr = np.stack([arr]*3, axis=-1)
-        return Image.fromarray(arr)
-    else:
-        return Image.open(uploaded_file).convert("RGB")
-
-# ─── STREAMLIT UI ───────────────────────────────────────────────────────────────
 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/JPEG/DCM) and click Generate Report.</p>",
-            unsafe_allow_html=True)
+st.markdown("<p style='text-align:center;'>Upload a chest X-ray and click Generate Report.</p>", unsafe_allow_html=True)
 
-if "stage" not in st.session_state:
-    st.session_state.stage = "upload"
-
-if st.session_state.stage == "upload":
-    uploaded = st.file_uploader(
-        "📤 Upload your chest X-ray",
-        type=["png","jpg","jpeg","dcm"],
-        label_visibility="visible"
-    )
+# File upload handling
+if "img" not in st.session_state:
+    uploaded = st.file_uploader("📤 Upload X-ray (PNG/JPG)", type=["png", "jpg", "jpeg"])
     if uploaded:
-        st.image(load_image(uploaded), width=350,
-                 caption=f"{uploaded.name} — {uploaded.size/1e6:.2f} MB")
-        if st.button("▶️ Generate Report"):
-            st.session_state.uploaded = uploaded
-            st.session_state.stage = "report"
-            st.experimental_rerun()
-
-elif st.session_state.stage == "report":
-    uploaded = st.session_state.uploaded
-    img = load_image(uploaded)
-
-    with st.spinner("🔍 Analyzing…"):
-        # 1) feature extraction
-        pixel_values = feat_ext(images=img, return_tensors="pt").pixel_values.to(device)
-        # 2) generation
-        output_ids = model.generate(
-            pixel_values,
-            max_length=64,
-            num_beams=4,
-            no_repeat_ngram_size=2,
-            early_stopping=True,
-        )
-        report = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
-    col1, col2 = st.columns(2)
-    with col1:
-        st.subheader("Your Uploaded X-ray")
-        st.image(img, use_column_width=True)
-        st.markdown(f"**File:** {uploaded.name}  \n**Size:** {uploaded.size/1e6:.2f} MB")
-    with col2:
-        st.subheader("📝 AI Diagnosis & Report")
-        st.markdown(
-            f"<div style='background:#e0f7fa;padding:12px;border-radius:6px;'>{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>
-<p style='text-align:center;color:gray;font-size:0.8em;'>
-  Powered by your fine-tuned ViT→T5 pipeline on Hugging Face.
-</p>
-""", unsafe_allow_html=True)
+        st.session_state.img = uploaded
+        st.experimental_rerun()
+    else:
+        st.stop()
+
+img_file = st.session_state.img
+img = Image.open(img_file).convert("RGB")
+st.image(img, use_column_width=True)
+
+col1, col2 = st.columns(2)
+with col1:
+    if st.button("▶️ Generate Report", use_container_width=True):
+        with st.spinner("Analyzing X-ray..."):
+            px = transform(img).unsqueeze(0).to(device)
+            out_ids = model.generate(px, max_length=128)
+            report = tokenizer.decode(out_ids[0], skip_special_tokens=True)
+        
+        st.subheader("📝 AI-Generated Report")
+        st.success(report)
+
+with col2:
+    if st.button("⬅️ Upload Another", use_container_width=True):
+        del st.session_state.img
+        st.experimental_rerun()