rollback to benchmark

cerebAI.py

@@ -9,43 +9,36 @@ from captum.attr import IntegratedGradients
 from typing import Tuple, Optional
 import albumentations as A
 from albumentations.pytorch import ToTensorV2
-import os 
-import requests 
-import pydicom 
-import io 
-import gc 
+import os
+import requests
+import pydicom  # REQUIRED FOR DICOM SUPPORT
+import io
+import gc  # For memory management
 
 # --- CONFIGURATION ---
-HF_MODEL_URL = "https://huggingface.co/arshenoy/cerebAI-stroke-model/resolve/main/best_model.pth" 
+HF_MODEL_URL = "https://huggingface.co/arshenoy/cerebAI-stroke-model/resolve/main/best_model.pth"
 DOWNLOAD_MODEL_PATH = "best_model_cache.pth"
 CLASS_LABELS = ['No Stroke', 'Ischemic Stroke', 'Hemorrhagic Stroke']
 IMAGE_SIZE = 224
-DEVICE = torch.device("cpu") 
-
-# --- INITIALIZE SESSION STATE (NEW: for stable UI) ---
-if 'results_ready' not in st.session_state:
-    st.session_state.results_ready = False
-    st.session_state.results = {}
-    st.session_state.file_bytes = None
-
+DEVICE = torch.device("cpu")  # For Streamlit Cloud stability
 
 # --- MODEL LOADING ---
 @st.cache_resource
 def load_model(model_url, local_path):
-    # ... (Model loading logic remains the same) ...
+    """Downloads model from URL if not cached, and loads the weights."""
     if not os.path.exists(local_path):
         st.info(f"Model not found locally. Downloading from remote repository...")
         try:
             response = requests.get(model_url, stream=True)
-            response.raise_for_status() 
+            response.raise_for_status()
+
             with open(local_path, "wb") as f:
                 for chunk in response.iter_content(chunk_size=8192):
                     f.write(chunk)
             st.success("Model download complete!")
-        except Exception:
-            st.error(f"FATAL ERROR: Could not download model.")
+        except Exception as e:
+            st.error(f"FATAL ERROR: Could not download model. Check the URL. Error: {e}")
             return None
-
     try:
         model = timm.create_model('convnext_base', pretrained=False)
         model.reset_classifier(num_classes=len(CLASS_LABELS))
@@ -53,19 +46,18 @@ def load_model(model_url, local_path):
         model.to(DEVICE)
         model.eval()
         return model
-    except Exception:
-        st.error(f"Failed to load model weights from cache.")
+    except Exception as e:
+        st.error(f"Failed to load model weights from cache. Error: {e}")
         return None
 
-# --- CORE HELPER FUNCTIONS (UNCHANGED) ---
-
+# --- HELPER FUNCTIONS ---
 def denormalize_image(tensor: torch.Tensor) -> np.ndarray:
     """Denormalizes a PyTorch tensor for matplotlib visualization."""
     if tensor.ndim == 4:
-        tensor = tensor.squeeze(0).detach() 
+        tensor = tensor.squeeze(0).detach()
     else:
-        tensor = tensor.detach() 
-    
+        tensor = tensor.detach()
+
     mean, std = np.array([0.5, 0.5, 0.5]), np.array([0.5, 0.5, 0.5])
     img = tensor.cpu().permute(1, 2, 0).numpy()
     img = (img * std) + mean
@@ -73,117 +65,81 @@ def denormalize_image(tensor: torch.Tensor) -> np.ndarray:
 
 def preprocess_image(image_bytes: bytes, file_name: str) -> Tuple[Optional[torch.Tensor], Optional[np.ndarray]]:
     """Loads, processes, and normalizes image, handling DICOM or JPG/PNG."""
-    
+    # 1. READ IMAGE DATA (Handles DICOM vs Standard formats)
     if file_name.lower().endswith(('.dcm', '.dicom')):
         try:
             dcm = pydicom.dcmread(io.BytesIO(image_bytes))
+
+            # FIX: Convert to Hounsfield Units (HU)
             pixel_array = dcm.pixel_array.astype(np.int16)
             slope = dcm.RescaleSlope
             intercept = dcm.RescaleIntercept
             pixel_array = pixel_array * slope + intercept
-            
-            window_center = 40 
+
+            # Apply Standard Brain Window (-100 HU to 150 HU)
+            window_center = 40
             window_width = 150
             min_hu = window_center - (window_width / 2)
             max_hu = window_center + (window_width / 2)
-            
+
+            # Apply the windowing transformation and scale to 0-255
             pixel_array[pixel_array < min_hu] = min_hu
             pixel_array[pixel_array > max_hu] = max_hu
             image_grayscale = ((pixel_array - min_hu) / (max_hu - min_hu) * 255).astype(np.uint8)
-            
         except Exception:
             return None, None
     else:
+        # Read standard image (PNG/JPG)
         image_grayscale = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_GRAYSCALE)
-        if image_grayscale is None: return None, None
-            
+        if image_grayscale is None:
+            return None, None
+
+    # 2. STANDARD PREPROCESSING
     image_rgb = cv2.cvtColor(cv2.resize(image_grayscale, (IMAGE_SIZE, IMAGE_SIZE)), cv2.COLOR_GRAY2RGB)
     image_norm = (image_rgb.astype(np.float32) / 255.0 - 0.5) / 0.5
     input_tensor = torch.tensor(image_norm, dtype=torch.float).permute(2, 0, 1).unsqueeze(0)
-    
     return input_tensor.to(DEVICE), image_rgb
 
 def generate_attribution(model: nn.Module, input_tensor: torch.Tensor, predicted_class_idx: int, n_steps: int = 20) -> np.ndarray:
     """Computes Integrated Gradients for the given input and class."""
-    target_class_int = int(predicted_class_idx) 
-    input_tensor.requires_grad_(True) 
-    
+    target_class_int = int(predicted_class_idx)
+
+    input_tensor.requires_grad_(True)
+
     ig = IntegratedGradients(model)
     baseline = torch.zeros_like(input_tensor).to(DEVICE)
-
     attributions_ig = ig.attribute(
         inputs=input_tensor,
         baselines=baseline,
         target=target_class_int,
         n_steps=n_steps
     )
-    
+
     attributions_ig_vis = attributions_ig.squeeze(0).sum(dim=0).abs().cpu().detach().numpy()
-    
     if attributions_ig_vis.max() > 0:
         attributions_ig_vis = attributions_ig_vis / attributions_ig_vis.max()
-    
     return attributions_ig_vis
 
 def plot_heatmap_and_original(original_image: np.ndarray, heatmap: np.ndarray, predicted_label: str):
     """Creates a Matplotlib figure for visualization."""
-    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5)) 
-    original_image_vis = (original_image.astype(np.float32) / 255.0) 
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))
 
+    original_image_vis = (original_image.astype(np.float32) / 255.0)
     ax1.imshow(original_image_vis)
     ax1.set_title("Original CT Scan", fontsize=14)
     ax1.axis('off')
 
     ax2.imshow(original_image_vis)
-    alpha_mask = heatmap * 0.7 + 0.3 
-
+    alpha_mask = heatmap * 0.7 + 0.3
     ax2.imshow(heatmap, cmap='jet', alpha=alpha_mask, vmin=0, vmax=1)
     ax2.set_title(f"Interpretation: {predicted_label}", fontsize=14)
     ax2.axis('off')
-    
+
     plt.tight_layout()
     return fig
 
-# -------------------- CORE DIAGNOSIS FUNCTION (TO BE CALLED BY BUTTON) --------------------
-
-def perform_full_analysis(model, image_bytes, file_name, n_steps_slider):
-    """Function called by the button to perform heavy computation and update state."""
-    
-    # 1. PREPARE INPUTS
-    input_tensor, original_image_rgb = preprocess_image(image_bytes, file_name) 
-
-    if input_tensor is None:
-        st.error("Could not process file. Ensure it is a valid DICOM/PNG/JPG.")
-        st.session_state.results_ready = False
-        return
-
-    # 2. PREDICT
-    with torch.no_grad():
-        output = model(input_tensor)
-        probabilities = torch.softmax(output, dim=1).squeeze(0).cpu().numpy()
-        predicted_class_idx = np.argmax(probabilities)
-    
-    # 3. GENERATE ATTRIBUTION (HEAVY PART)
-    heatmap = generate_attribution(model, input_tensor, predicted_class_idx, n_steps=n_steps_slider)
-    
-    # 4. CRITICAL MEMORY MANAGEMENT
-    del input_tensor 
-    del output
-    gc.collect() 
-
-    # 5. SAVE FINAL RESULTS TO STATE
-    st.session_state.results = {
-        'label': CLASS_LABELS[predicted_class_idx],
-        'confidence': probabilities[predicted_class_idx],
-        'probabilities': probabilities,
-        'image_rgb': original_image_rgb,
-        'heatmap': heatmap
-    }
-    st.session_state.results_ready = True
-    st.rerun() # Force a single, clean render of the results
-
-
 # ==============================================================================
+
 # -------------------- STREAMLIT FRONTEND --------------------
 # ==============================================================================
 
@@ -191,76 +147,78 @@ st.set_page_config(page_title="CerebAI: Stroke Prediction Dashboard", layout="wi
 st.title("CerebAI: AI-Powered Stroke Detection")
 st.markdown("---")
 
+# FIX: Load the model using the download mechanism
 model = load_model(HF_MODEL_URL, DOWNLOAD_MODEL_PATH)
 
 if model is not None:
-    # --- INPUT CONTROLS ---
+    # --- INTERACTIVE CONTROLS (Sidebar or Main Area) ---
     st.markdown("### Analysis Controls")
-    
     n_steps_slider = st.slider(
         'Integration Steps (Affects Accuracy & Speed)',
-        min_value=5, 
-        max_value=50, 
-        value=10, 
+        min_value=5,
+        max_value=50,
+        value=20,
         step=5,
-        key="n_steps_slider_key", # Added key to prevent slider reset on rerun
         help="Higher steps (up to 50) provide a smoother, more accurate heatmap but use more CPU."
     )
     st.markdown("---")
 
-
     # --- FILE UPLOAD ---
     st.markdown("### Upload CT Scan Image")
     uploaded_file = st.file_uploader(
-        "Choose a Dicom, PNG, JPG, or JPEG file", 
-        type=["dcm", "dicom", "png", "jpg", "jpeg"],
-        key="file_uploader" # Added key to manage file state
+        "Choose a Dicom, PNG, JPG, or JPEG file",
+        type=["dcm", "dicom", "png", "jpg", "jpeg"]
     )
-    
-    
-    # --- UI FLOW MANAGEMENT ---
+
     if uploaded_file is not None:
         image_bytes = uploaded_file.read()
-        file_name = uploaded_file.name 
-
-        # Display the image immediately (fast operation, no lag)
-        st.subheader("Input Image")
-        st.image(image_bytes, use_container_width=True, caption=file_name) 
-
-        # --- GATED BUTTON ---
-        if st.button("▶️ RUN DIAGNOSIS & ATTRIBUTION", type="primary", use_container_width=True, key="run_button"):
-            with st.spinner(f'Calculating Integrated Gradients ({n_steps_slider} steps)...'):
-                # Execute the heavy logic and save to session state
-                perform_full_analysis(model, image_bytes, file_name, n_steps_slider)
-                
-    # --- RESULTS DISPLAY (ONLY RUNS AFTER BUTTON CLICK) ---
-    if st.session_state.results_ready:
-        data = st.session_state.results
-        
-        st.markdown("---")
-        st.markdown("### 2. Diagnosis and Interpretation")
-
-        col_metric, col_plot = st.columns([1, 2])
-
-        with col_metric:
-            st.subheader("Prediction Summary")
-            st.metric(
-                label="Diagnosis", 
-                value=data['label'],
-                delta=f"{data['confidence']*100:.2f}% Confidence",
-                delta_color='normal' 
-            )
-            st.markdown("---")
-            st.subheader("Confidence Breakdown")
-            prob_data = {
-                'Class': CLASS_LABELS,
-                'Confidence': [f"{p:.4f}" for p in data['probabilities']]
-            }
-            st.dataframe(prob_data, hide_index=True, use_container_width=True)
-
-        with col_plot:
-            st.subheader("Visual Explanation")
-            fig = plot_heatmap_and_original(data['image_rgb'], data['heatmap'], data['label'])
-            st.pyplot(fig, clear_figure=True, use_container_width=True) 
-
-        st.success("Analysis Complete: The heatmap highlights the regions most critical to the diagnosis.")
+        file_name = uploaded_file.name
+
+        # 1. PROCESS IMAGE FIRST (Defines original_image_rgb)
+        input_tensor, original_image_rgb = preprocess_image(image_bytes, file_name)
+
+        # --- DISPLAY AND RESULTS LAYOUT ---
+        col1, col2 = st.columns(2)
+
+        with col1:
+            st.subheader("Uploaded Image")
+            # Display the processed NumPy array
+            st.image(original_image_rgb, use_container_width=True, caption=file_name)
+
+        # Run Prediction and Attribution
+        if input_tensor is not None:
+            # Predict
+            with torch.no_grad():
+                output = model(input_tensor)
+                probabilities = torch.softmax(output, dim=1).squeeze(0).cpu().numpy()
+                predicted_class_idx = np.argmax(probabilities)
+
+                predicted_label = CLASS_LABELS[predicted_class_idx]
+            confidence_score = probabilities[predicted_class_idx]
+
+            # Generate Attribution
+            heatmap = generate_attribution(model, input_tensor, predicted_class_idx, n_steps=n_steps_slider)
+
+            # CRITICAL MEMORY MANAGEMENT
+            del input_tensor
+            del output
+            gc.collect()
+
+            with col2:
+                st.subheader("Prediction Summary")
+
+                st.metric(
+                    label="Diagnosis",
+                    value=predicted_label,
+                    delta=f"{confidence_score*100:.2f}% Confidence",
+                    delta_color='normal'
+                )
+
+                st.markdown("---")
+                st.subheader("Confidence Breakdown")
+
+                prob_data = {
+                    'Class': CLASS_LABELS,
+                    'Confidence': [f"{p:.4f}" for p in probabilities]
+                }
+                st.dataframe(prob_data, hide_index=True, use_container_width