Update app.py

app.py

@@ -7,21 +7,26 @@ from collections import Counter
 import numpy as np
 from huggingface_hub import hf_hub_download
 from PIL import Image
+import torch.serialization
+from ultralytics.nn.tasks import DetectionModel
+
+# Allow YOLO model class when loading torch checkpoints
+torch.serialization.add_safe_globals([DetectionModel])
 
 # -------------------------------
-# Load model safely
+# Load YOLO model safely
 # -------------------------------
 def load_model():
-    model = None
     try:
         print("🔧 Attempting to load YOLO model...")
 
-        # Prefer local cache to avoid re-download
         model_path = os.path.join("models", "best.pt")
         os.makedirs("models", exist_ok=True)
 
+        # Check if model exists locally
         if not os.path.exists(model_path):
             print("🌐 Model not found locally, downloading from Hugging Face...")
+
             if "HF_TOKEN" in os.environ:
                 model_path = hf_hub_download(
                     repo_id="Faethon88/sar",
@@ -36,16 +41,16 @@ def load_model():
         else:
             print(f"📁 Using local model: {model_path}")
 
-        # Load YOLO model (this handles custom weights automatically)
+        # Load YOLO model
         model = YOLO(model_path)
         print("✅ Model loaded successfully!")
         print(f"📋 Classes: {list(model.names.values())}")
+        return model
 
     except Exception as e:
         print(f"❌ Model load failed: {e}")
-        model = None
-
-    return model
+        print("⚠️ Running in DEMO mode without model...")
+        return None
 
 
 # -------------------------------
@@ -60,9 +65,9 @@ print(f"📊 Model loaded: {MODEL_LOADED}")
 # Detection function
 # -------------------------------
 def detect_ships(image_input, confidence=0.5):
-    """Detect ships in SAR image — supports both demo and real mode."""
+    """Detect ships in SAR image — works in real or demo mode."""
     try:
-        # Load input image properly
+        # Load image from file path or numpy
         if isinstance(image_input, str):
             print(f"📁 Loading image from path: {image_input}")
             image = Image.open(image_input).convert("RGB")
@@ -72,45 +77,45 @@ def detect_ships(image_input, confidence=0.5):
         else:
             image_np = np.array(image_input)
 
-        # Convert to BGR for OpenCV drawing
+        # Convert to BGR for OpenCV
         if len(image_np.shape) == 3 and image_np.shape[2] == 3:
             image_np = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
 
         # DEMO MODE
         if not MODEL_LOADED:
-            print("🎭 Running in DEMO mode (no model loaded)")
-            annotated_image = image_np.copy()
-            h, w = annotated_image.shape[:2]
-            cv2.rectangle(annotated_image, (w//4, h//4), (w//2, h//2), (0, 255, 0), 3)
-            cv2.putText(annotated_image, "ship 0.85", (w//4, h//4 - 10),
+            print("🎭 Running in DEMO mode")
+            annotated = image_np.copy()
+            h, w = annotated.shape[:2]
+            cv2.rectangle(annotated, (w//4, h//4), (w//2, h//2), (0, 255, 0), 3)
+            cv2.putText(annotated, "ship 0.85", (w//4, h//4 - 10),
                         cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
-            summary = f"Demo Results (Model not loaded):\n- ship: 1\n\nConfidence: {confidence:.2f}\nTotal: 1"
-            annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
-            return annotated_image, summary
+            annotated = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
+            summary = "Demo Mode:\n- ship: 1\n\nConfidence: 0.85\nTotal: 1"
+            return annotated, summary
 
         # REAL MODEL PREDICTION
         print(f"🔍 Running real model (conf={confidence})...")
         results = model.predict(image_np, conf=confidence, verbose=False)
         result = results[0]
 
-        annotated_image = image_np.copy()
-        boxes = result.boxes.xyxy.cpu().numpy() if result.boxes is not None else []
-        confs = result.boxes.conf.cpu().numpy().tolist() if result.boxes is not None else []
-        class_ids = result.boxes.cls.cpu().numpy().tolist() if result.boxes is not None else []
+        annotated = image_np.copy()
+        boxes = result.boxes.xyxy.cpu().numpy() if result.boxes else []
+        confs = result.boxes.conf.cpu().numpy().tolist() if result.boxes else []
+        class_ids = result.boxes.cls.cpu().numpy().tolist() if result.boxes else []
 
-        # Draw detections
         class_names = []
         for box, cls_id, conf in zip(boxes, class_ids, confs):
             x1, y1, x2, y2 = map(int, box)
             cls_name = model.names.get(int(cls_id), "unknown")
             class_names.append(cls_name)
             label = f"{cls_name} {conf:.2f}"
-            cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 3)
-            cv2.putText(annotated_image, label, (x1, y1 - 10),
+            cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 255, 0), 3)
+            cv2.putText(annotated, label, (x1, y1 - 10),
                         cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)
 
-        annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
+        annotated = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
         counts = Counter(class_names)
+
         summary = "Detection Results:\n"
         if counts:
             summary += "\n".join([f"- {cls}: {cnt}" for cls, cnt in counts.items()])
@@ -120,11 +125,11 @@ def detect_ships(image_input, confidence=0.5):
         summary += f"\nTotal detections: {len(class_names)}"
 
         print(f"✅ Detected {len(class_names)} objects.")
-        return annotated_image, summary
+        return annotated, summary
 
     except Exception as e:
         print(f"❌ Error during detection: {e}")
-        return None, f"Processing error: {str(e)}"
+        return None, f"Processing error: {e}"
 
 
 # -------------------------------
@@ -143,18 +148,11 @@ demo = gr.Interface(
     title="🛰️ SAR Ship Detection",
     description="Upload a SAR image to detect ships"
                 + (" (DEMO MODE - Model not loaded)" if not MODEL_LOADED else ""),
-    examples=[
-        ["https://upload.wikimedia.org/wikipedia/commons/6/6d/Sentinel-1_SAR_image_of_ship_traffic.jpg", 0.5]
-    ] if not MODEL_LOADED else None
 )
 
 # -------------------------------
-# Main
+# Main Entry Point
 # -------------------------------
 if __name__ == "__main__":
     print("🚀 Starting SAR Ship Detection App...")
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=True
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)