Update app.py

app.py

@@ -9,12 +9,12 @@ import albumentations as A
 from albumentations.pytorch import ToTensorV2
 from transformers import AutoModel
 import matplotlib.pyplot as plt
+import spaces
 
 # ==============================================================================
 # 1. GLOBAL CONFIGURATION & PATHOLOGY DEFINITIONS
 # ==============================================================================
 IMAGE_SIZE = 384
-DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 PATHOLOGY_COLS = [
     'Enlarged Cardiomediastinum', 'Cardiomegaly', 'Lung Opacity',
@@ -154,21 +154,17 @@ class ViTGradCAM:
 # ==============================================================================
 # 3. INITIALIZE MODEL & WEIGHTS
 # ==============================================================================
-def load_trained_model():
-    model = OptimizedMultiModalRadDINO(num_classes=13).to(DEVICE)
+def load_trained_model(device):
+    model = OptimizedMultiModalRadDINO(num_classes=13).to(device)
     model_path = 'best_model.pt'
     if os.path.exists(model_path):
-        model.load_state_dict(torch.load(model_path, map_location=DEVICE))
+        model.load_state_dict(torch.load(model_path, map_location=device))
         print("Successfully loaded trained weights from best_model.pt")
     else:
         print("Warning: best_model.pt not found. Initializing with pretrained backbone.")
     model.eval()
     return model
 
-model = load_trained_model()
-target_layer = model.vision_model.encoder.layer[-1].norm1
-cam_engine = ViTGradCAM(model, target_layer)
-
 # ==============================================================================
 # 4. PREPROCESSING PIPELINE
 # ==============================================================================
@@ -186,12 +182,18 @@ val_transforms = A.Compose([
 ])
 
 # ==============================================================================
-# 5. INFERENCE & VISUALIZATION LOGIC
+# 5. INFERENCE & VISUALIZATION LOGIC (DECORATED WITH @spaces.GPU)
 # ==============================================================================
+@spaces.GPU
 def predict_xray(image_input, view_position, target_cam_pathology):
     if image_input is None:
         return None, None, "Please upload a valid chest X-ray image."
 
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = load_trained_model(device)
+    target_layer = model.vision_model.encoder.layer[-1].norm1
+    cam_engine = ViTGradCAM(model, target_layer)
+
     # Convert image format
     image_rgb = cv2.cvtColor(image_input, cv2.COLOR_BGR2RGB) if len(image_input.shape) == 3 else image_input
     
@@ -204,11 +206,11 @@ def predict_xray(image_input, view_position, target_cam_pathology):
     elif view_position == "LL":
         view_vector = [0.0, 0.0, 1.0]
         
-    tabular_tensor = torch.tensor([view_vector], dtype=torch.float32).to(DEVICE)
+    tabular_tensor = torch.tensor([view_vector], dtype=torch.float32).to(device)
     
     # Process Image Tensor
     augmented = val_transforms(image=image_rgb)
-    img_tensor = augmented['image'].unsqueeze(0).to(DEVICE)
+    img_tensor = augmented['image'].unsqueeze(0).to(device)
 
     # Run Model Inference
     with torch.no_grad():
@@ -312,7 +314,7 @@ with gr.Blocks(css=custom_css, title="ThoraxVision AI - Clinical X-Ray Analysis"
         # TAB 2: TECHNICAL DOCUMENTATION & ARCHITECTURE
         # ----------------------------------------------------------------------
         with gr.TabItem("📖 System Architecture & Scientific Documentation"):
-            gr.Markdown("""
+            gr.Markdown(r"""
             ## Executive Overview
             **ThoraxVision AI** is an advanced multimodal diagnostic deep learning framework trained to detect 13 distinct thoracic pathologies simultaneously from chest radiographs. Built upon the `microsoft/rad-dino` Vision Transformer (ViT) architecture, the model integrates clinical tabular metadata with high-resolution image representations.
 
@@ -356,7 +358,7 @@ with gr.Blocks(css=custom_css, title="ThoraxVision AI - Clinical X-Ray Analysis"
             metrics_df = pd.DataFrame(MODEL_METRICS)
             gr.Dataframe(value=metrics_df, interactive=False)
             
-            gr.Markdown("""
+            gr.Markdown(r"""
             * **Mean AUC-ROC Across All Classes:** **0.7852**
             * **Optimal Threshold Strategy:** Decision thresholds were tuned individually for each pathology to maximize the F1-Score rather than relying on an arbitrary 0.5 default cut-off.
             """)