Update app.py

app.py

@@ -1,5 +1,4 @@
 import gradio as gr
-import cv2
 import torch
 import torchvision
 import numpy as np
@@ -7,6 +6,7 @@ from PIL import Image
 import matplotlib.pyplot as plt
 import matplotlib.patches as patches
 from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+from huggingface_hub import hf_hub_download
 import os
 import io
 
@@ -16,12 +16,16 @@ os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
 CLASS_NAMES = {1: 'Nipple', 2: 'Lump'}
 CLASS_COLORS = {1: 'white', 2: 'white'}
 
+# Global variables for model (load once at startup)
+MODEL = None
+DEVICE = None
+
 def preprocess_image(image):
     """Load and preprocess image for Faster R-CNN."""
     # Convert PIL Image to numpy array
     image = np.array(image)
     
-    # Convert RGB to RGB (already in correct format from Gradio)
+    # Already in RGB format from Gradio
     image = image.astype(np.float32) / 255.0  # Normalize to [0,1]
     
     # Normalize using ImageNet mean and std
@@ -40,13 +44,41 @@ def load_model(checkpoint_path, device):
     model.to(device).eval()
     return model
 
+def initialize_model():
+    """Initialize model at startup by downloading from HF Hub."""
+    global MODEL, DEVICE
+    
+    if MODEL is None:
+        DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        print(f"Downloading model from IFMedTech/Lumps repository...")
+        # Download model from private HuggingFace repository
+        # Token is automatically read from HF_TOKEN environment variable (Spaces secrets)
+        try:
+            checkpoint_path = hf_hub_download(
+                repo_id="IFMedTech/Lumps",
+                filename="lumps.pth",
+                repo_type="model",
+                token=os.environ.get("HF_TOKEN")  # Use token from Spaces secrets
+            )
+            print(f"Model downloaded to: {checkpoint_path}")
+            print(f"Loading model on {DEVICE}...")
+            MODEL = load_model(checkpoint_path, DEVICE)
+            print(f"Model loaded successfully!")
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            raise RuntimeError(
+                f"Failed to load model from IFMedTech/Lumps. "
+                f"Please ensure HF_TOKEN is set in Spaces secrets with read access to the private repository. "
+                f"Error: {e}"
+            )
+    
+    return MODEL, DEVICE
+
 def predict(image, score_thresh=0.5):
     """Run inference and return image with bounding boxes."""
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    
-    # Load model
-    checkpoint_path = "lumps.pth"  # This will be downloaded from the model repo
-    model = load_model(checkpoint_path, device)
+    # Ensure model is loaded
+    model, device = initialize_model()
     
     # Preprocess image
     image_tensor = preprocess_image(image)