Update app.py

app.py

@@ -6,6 +6,8 @@ from collections import Counter
 from PIL import Image
 from huggingface_hub import hf_hub_download
 import os
+import base64
+from io import BytesIO
 
 # -------------------------------
 # Load YOLO model safely with HF token
@@ -75,9 +77,28 @@ def detect_ships(image: Image.Image, confidence: float):
         return None, f"❌ Detection failed: {e}"
 
 # -------------------------------
-# Gradio API function
+# Gradio API function with wrapper for remote dict input
 # -------------------------------
 def predict(image, confidence):
+    print("DEBUG: predict called")
+    print("DEBUG: raw image type:", type(image))
+    print("DEBUG: confidence type/value:", type(confidence), confidence)
+
+    # Handle dict input from remote client (Flask sends {"name":..., "data": data_uri})
+    if isinstance(image, dict):
+        data = image.get("data") or image.get("image") or ""
+        if data and isinstance(data, str) and data.startswith("data:image"):
+            try:
+                header, b64 = data.split(",", 1)
+                image = Image.open(BytesIO(base64.b64decode(b64))).convert("RGB")
+                print("DEBUG: decoded image from data URI to PIL.Image")
+            except Exception as e:
+                print("ERROR decoding data URI:", e)
+                raise
+
+    if image is None:
+        raise ValueError("Received empty image")
+
     return detect_ships(image, confidence)
 
 # -------------------------------
@@ -103,4 +124,4 @@ if __name__ == "__main__":
         server_port=7860,
         show_error=True,  # Show detailed errors in browser
         debug=True        # Print detailed logs to console
-    )
+    )