Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,7 +3,8 @@
 # ──── SET ENVIRONMENT VARIABLES BEFORE ANY IMPORTS ──────────────────────────────
 import os
 import tempfile
-import requests
+import sys
+import atexit
 
 # Create dedicated cache directories
 CACHE_DIR = "/tmp/hf_cache"
@@ -21,7 +22,8 @@ os.environ.update({
     "TRANSFORMERS_CACHE": f"{CACHE_DIR}/transformers",
     "HF_HUB_CACHE": f"{CACHE_DIR}/huggingface_hub",
     "HUGGINGFACE_HUB_CACHE": f"{CACHE_DIR}/huggingface_hub",
-    "HF_HUB_DISABLE_TELEMETRY": "1"  # Disable telemetry to reduce rate limiting
+    "HF_HUB_DISABLE_TELEMETRY": "1",
+    "STREAMLIT_SERVER_ENABLE_FILE_WATCHER": "false"
 })
 
 # Create all cache directories explicitly
@@ -39,6 +41,17 @@ if not os.path.exists(CONFIG_TOML):
     with open(CONFIG_TOML, "w") as f:
         f.write("[browser]\n")
         f.write("gatherUsageStats = false\n")
+        f.write("[server]\n")
+        f.write("fileWatcherType = none\n")
+
+# Monkey-patch Streamlit to prevent root directory access
+def safe_makedirs(name, mode=0o777, exist_ok=False):
+    """Prevent Streamlit from creating directories outside /tmp"""
+    if name.startswith(('/.streamlit', '/root', '/home')):
+        name = name.replace('/', '/tmp/', 1)
+    return os.makedirs(name, mode, exist_ok)
+
+os.makedirs = safe_makedirs
 
 # ──── NOW IMPORT OTHER LIBRARIES ───────────────────────────────────────────────
 import json
@@ -52,9 +65,7 @@ from huggingface_hub import hf_hub_download, HfApi
 
 # ──── MODEL DEFINITION ─────────────────────────────────────────────────────────
 MODEL_ID = "RakeshNJ12345/Chest-Radiology"
-
-# Alternative model access through proxy
-PROXY_URL = "https://hf-mirror.com"
+PROXY_URL = "https://hf-mirror.com"  # Proxy for Hugging Face downloads
 
 class TwoViewVisionReportModel(nn.Module):
     def __init__(self, vit: ViTModel, t5: T5ForConditionalGeneration, tokenizer: T5Tokenizer):
@@ -65,7 +76,7 @@ class TwoViewVisionReportModel(nn.Module):
         self.tokenizer = tokenizer
         self.t5 = t5
 
-    def generate(self, img: torch.Tensor, max_length: int = 64) -> torch.Tensor:
+    def generate(self, img: torch.Tensor, max_length: int = 128) -> torch.Tensor:
         device = img.device
         vf = self.vit(pixel_values=img).pooler_output
         pf = self.proj_f(vf).unsqueeze(1)
@@ -95,7 +106,7 @@ class TwoViewVisionReportModel(nn.Module):
         )
         return out_ids
 
-# ──── MODEL LOADING WITH PROXY SUPPORT AND ERROR HANDLING ──────────────────────
+# ──── MODEL LOADING WITH ERROR HANDLING ────────────────────────────────────────
 @st.cache_resource(show_spinner=False)
 def load_models():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -108,7 +119,6 @@ def load_models():
     ]:
         os.makedirs(path, exist_ok=True)
 
-    # Try to download using standard method first
     try:
         # Download config
         cfg_path = hf_hub_download(
@@ -119,20 +129,17 @@ def load_models():
             local_files_only=False
         )
     except Exception as e:
-        st.error(f"❌ Failed to download model via Hugging Face Hub: {str(e)}")
+        st.error(f"❌ Failed to download model config: {str(e)}")
         st.info("⚠️ Trying alternative download method...")
-        
-        # Use proxy mirror
-        cfg_path = f"{CACHE_DIR}/huggingface_hub/config.json"
         api = HfApi(endpoint=PROXY_URL)
-        api.hf_hub_download(
+        cfg_path = api.hf_hub_download(
             repo_id=MODEL_ID,
             filename="config.json",
             repo_type="model",
             cache_dir=f"{CACHE_DIR}/huggingface_hub",
             local_files_only=False
         )
-
+    
     cfg = json.load(open(cfg_path, "r"))
 
     # Load models with explicit cache directories
@@ -142,8 +149,8 @@ def load_models():
             ignore_mismatched_sizes=True,
             cache_dir=f"{CACHE_DIR}/transformers"
         ).to(device)
-    except:
-        # Use proxy if standard download fails
+    except Exception as e:
+        st.warning(f"⚠️ Standard ViT download failed: {str(e)}")
         vit = ViTModel.from_pretrained(
             "google/vit-base-patch16-224",
             ignore_mismatched_sizes=True,
@@ -156,7 +163,8 @@ def load_models():
             "t5-base",
             cache_dir=f"{CACHE_DIR}/transformers"
         ).to(device)
-    except:
+    except Exception as e:
+        st.warning(f"⚠️ Standard T5 download failed: {str(e)}")
         t5 = T5ForConditionalGeneration.from_pretrained(
             "t5-base",
             cache_dir=f"{CACHE_DIR}/transformers",
@@ -168,7 +176,8 @@ def load_models():
             MODEL_ID,
             cache_dir=f"{CACHE_DIR}/transformers"
         )
-    except:
+    except Exception as e:
+        st.warning(f"⚠️ Standard tokenizer download failed: {str(e)}")
         tok = T5Tokenizer.from_pretrained(
             MODEL_ID,
             cache_dir=f"{CACHE_DIR}/transformers",
@@ -186,8 +195,8 @@ def load_models():
             cache_dir=f"{CACHE_DIR}/huggingface_hub",
             local_files_only=False
         )
-    except:
-        # Use proxy mirror for model weights
+    except Exception as e:
+        st.warning(f"⚠️ Standard model weights download failed: {str(e)}")
         api = HfApi(endpoint=PROXY_URL)
         ckpt_path = api.hf_hub_download(
             repo_id=MODEL_ID,
@@ -201,11 +210,12 @@ def load_models():
     model.load_state_dict(state)
     return device, model, tok
 
-# ──── APP INTERFACE WITH ERROR HANDLING ───────────────────────────────────────
+# ──── APP INTERFACE ───────────────────────────────────────────────────────────
 try:
     device, model, tokenizer = load_models()
 except Exception as e:
     st.error(f"🚨 Critical Error: Failed to load models. {str(e)}")
+    st.info("Please try refreshing the page or contact support@example.com")
     st.stop()
 
 transform = T.Compose([
@@ -214,7 +224,24 @@ transform = T.Compose([
     T.Normalize(mean=0.5, std=0.5),
 ])
 
-st.set_page_config(page_title="Radiology Report Analysis", layout="wide")
+st.set_page_config(
+    page_title="Radiology Report Analysis",
+    layout="wide",
+    # Disable Streamlit's default behavior that causes permission issues
+    initial_sidebar_state="collapsed"
+)
+
+# Custom CSS to hide Streamlit elements that might cause issues
+st.markdown("""
+<style>
+    .reportview-container .main .block-container {padding-top: 2rem;}
+    header {visibility: hidden;}
+    .stDeployButton {display:none;}
+    #MainMenu {visibility: hidden;}
+    footer {visibility: hidden;}
+</style>
+""", unsafe_allow_html=True)
+
 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 and click Generate Report.</p>", unsafe_allow_html=True)
 
@@ -223,9 +250,9 @@ if "img" not in st.session_state:
     uploaded = st.file_uploader("📤 Upload X-ray (PNG/JPG)", type=["png", "jpg", "jpeg"])
     if uploaded:
         try:
-            # Validate image
             img = Image.open(uploaded).convert("RGB")
-            img.verify()  # Check if image is valid
+            # Quick verification by thumbnail generation
+            img.thumbnail((10, 10))
             st.session_state.img = uploaded
             st.experimental_rerun()
         except Exception as e:
@@ -240,8 +267,8 @@ 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..."):
+    if st.button("▶️ Generate Report", use_container_width=True, type="primary"):
+        with st.spinner("Analyzing X-ray. This may take 10-20 seconds..."):
             try:
                 px = transform(img).unsqueeze(0).to(device)
                 out_ids = model.generate(px, max_length=128)
@@ -261,8 +288,8 @@ with col2:
 # Add footer with troubleshooting
 st.markdown("---")
 st.markdown("""
-**Troubleshooting Tips:**
-- If model download fails, wait 5 minutes and refresh
-- Use standard chest X-ray images in PNG or JPG format
-- For persistent errors, contact support@example.com
+**Note:** 
+- First-time model loading may take 1-2 minutes
+- For optimal results, use clear chest X-ray images
+- Contact support@example.com for assistance
 """)