Initial implementation of MosqScope

app.py

@@ -1,117 +1,138 @@
 import torch
 import torchvision.transforms as transforms
-from torchvision.models.detection import ssd300_vgg16
+from torchvision.models.detection.ssd import ssd300_vgg16
 import av
 import numpy as np
 import cv2
 import streamlit as st
 from streamlit_webrtc import webrtc_streamer, VideoProcessorBase, WebRtcMode, RTCConfiguration
 from huggingface_hub import hf_hub_download
+import asyncio
 import logging
+import os
 
-# Set up logging
-logging.basicConfig(level=logging.DEBUG)
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+
+# Fix for asyncio loop issues in some environments
+os.environ["STREAMLIT_SERVER_ENABLE_STATIC_SERVING"] = "true"
 
 # Define dataset classes
 classes = ['dengue-regions', 'wet_surface']
-num_classes = len(classes) + 1  # Including background
 
 # Load the SSD Model
 @st.cache_resource
 def load_model():
     try:
         model_path = hf_hub_download(repo_id="DhominickJ/MosqScope", filename="mosquito_model.pth")
-        model = ssd300_vgg16(pretrained=False)  # Don't load ImageNet weights
-        
-        # SSD models have a different structure - no need to modify the head like in Faster R-CNN
+        model = ssd300_vgg16(pretrained=False)
         model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
         model.eval()
         return model
     except Exception as e:
         st.error(f"Error loading model: {str(e)}")
+        logging.error(f"Model loading error: {e}")
         return None
 
-try:
-    model = load_model()
-except Exception as e:
-    st.error(f"Error loading model: {e}")
-    model = None
-
-# Define Video Processor for WebRTC
-class SSDVideoProcessor(VideoProcessorBase):
-    def __init__(self):
-        self.model = model
-        # SSD models expect input in [0,1] range and resized to 300x300
-        self.transform = transforms.Compose([
-            transforms.ToPILImage(),
-            transforms.Resize((300, 300)),
-            transforms.ToTensor(),
-        ])
-
+# Simple fallback class if model loading fails
+class VideoProcessor(VideoProcessorBase):
     def recv(self, frame):
-        if self.model is None:
-            # Just return the frame if model isn't loaded
-            return frame
-            
         img = frame.to_ndarray(format="bgr24")
-        # Make a copy for drawing
-        display_img = img.copy()
-        
-        try:
-            # Transform for model
-            image_tensor = self.transform(img).unsqueeze(0)
-
-            with torch.no_grad():
-                detections = self.model(image_tensor)
-                
-            # Get the detection results
-            boxes = detections[0]['boxes'].cpu().numpy()
-            scores = detections[0]['scores'].cpu().numpy()
-            labels = detections[0]['labels'].cpu().numpy()
-                
-            # Scale coordinates to original image dimensions
-            h, w = img.shape[:2]
-            scale_x, scale_y = w / 300, h / 300
-                
-            # Draw detections
-            for box, label, score in zip(boxes, labels, scores):
-                if score > 0.5:  # Only show confident detections
-                    x_min, y_min, x_max, y_max = box
-                    # Scale coordinates back to original image
-                    x_min, x_max = int(x_min * scale_x), int(x_max * scale_x)
-                    y_min, y_max = int(y_min * scale_y), int(y_max * scale_y)
-                    
-                    cv2.rectangle(display_img, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
-                    label_name = classes[label - 1]  # Adjust for background class
-                    cv2.putText(display_img, f"{label_name} {score:.2f}", 
-                                (x_min, y_min - 5), cv2.FONT_HERSHEY_SIMPLEX, 
-                                0.5, (0, 0, 255), 2)
-        except Exception as e:
-            logging.error(f"Error in inference: {e}")
-            # Add error message to frame
-            cv2.putText(display_img, f"Error: {str(e)}", (10, 30), 
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
-
-        return av.VideoFrame.from_ndarray(display_img, format="bgr24")
+        # Just add a text overlay indicating the model isn't loaded
+        cv2.putText(img, "Model not loaded", (10, 30), 
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
+        return av.VideoFrame.from_ndarray(img, format="bgr24")
 
 # Streamlit UI
 st.title("Mosquito Detection with WebRTC")
 st.write("This app uses a SSD model to detect mosquito breeding sites in real-time.")
 
-# Configure WebRTC with proper STUN/TURN servers
-rtc_config = RTCConfiguration(
-    {"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]},
+# Use a more reliable WebRTC configuration
+rtc_configuration = RTCConfiguration(
+    {"iceServers": [
+        {"urls": ["stun:stun.l.google.com:19302"]},
+        {
+            "urls": ["turn:openrelay.metered.ca:80"],
+            "username": "openrelayproject",
+            "credential": "openrelayproject",
+        }
+    ]}
 )
 
-# Start WebRTC Streaming with proper error handling
+# Load model conditionally - separate from the WebRTC setup
+try:
+    model = load_model()
+    
+    if model is not None:
+        # Define Video Processor with the loaded model
+        class SSDVideoProcessor(VideoProcessorBase):
+            def __init__(self):
+                self.model = model
+                self.transform = transforms.Compose([
+                    transforms.ToPILImage(),
+                    transforms.Resize((300, 300)),
+                    transforms.ToTensor(),
+                ])
+
+            def recv(self, frame):
+                img = frame.to_ndarray(format="bgr24")
+                display_img = img.copy()
+                
+                try:
+                    # Transform for model
+                    image_tensor = self.transform(img).unsqueeze(0)
+
+                    with torch.no_grad():
+                        detections = self.model(image_tensor)
+                        
+                    # Get the detection results
+                    boxes = detections[0]['boxes'].cpu().numpy()
+                    scores = detections[0]['scores'].cpu().numpy()
+                    labels = detections[0]['labels'].cpu().numpy()
+                        
+                    # Scale coordinates to original image dimensions
+                    h, w = img.shape[:2]
+                    scale_x, scale_y = w / 300, h / 300
+                        
+                    # Draw detections
+                    for box, label, score in zip(boxes, labels, scores):
+                        if score > 0.5:  # Only show confident detections
+                            x_min, y_min, x_max, y_max = box
+                            # Scale coordinates back to original image
+                            x_min, x_max = int(x_min * scale_x), int(x_max * scale_x)
+                            y_min, y_max = int(y_min * scale_y), int(y_max * scale_y)
+                            
+                            cv2.rectangle(display_img, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
+                            label_name = classes[label - 1]  # Adjust for background class
+                            cv2.putText(display_img, f"{label_name} {score:.2f}", 
+                                        (x_min, y_min - 5), cv2.FONT_HERSHEY_SIMPLEX, 
+                                        0.5, (0, 0, 255), 2)
+                except Exception as e:
+                    logging.error(f"Error in inference: {e}")
+                    cv2.putText(display_img, f"Error: {str(e)}", (10, 30), 
+                                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
+
+                return av.VideoFrame.from_ndarray(display_img, format="bgr24")
+                
+        processor_factory = SSDVideoProcessor
+    else:
+        st.warning("Model couldn't be loaded. Running in fallback mode.")
+        processor_factory = VideoProcessor
+        
+except Exception as e:
+    st.error(f"Error setting up model: {e}")
+    processor_factory = VideoProcessor
+
+# Start WebRTC streaming in a try-except block
 try:
+    # Use simpler configuration with fewer options to reduce chances of error
     webrtc_ctx = webrtc_streamer(
-        key="ssd-detection",
+        key="mosquito-detection",
         mode=WebRtcMode.SENDRECV,
-        rtc_configuration=rtc_config,
-        video_processor_factory=SSDVideoProcessor,
+        rtc_configuration=rtc_configuration,
+        video_processor_factory=processor_factory,
         media_stream_constraints={"video": True, "audio": False},
-        async_processing=True,
+        async_processing=False,  # Try with sync processing first
     )
 except Exception as e:
     st.error(f"WebRTC Error: {e}")