Initial implementation of MosqScope

app.py

@@ -5,8 +5,12 @@ import av
 import numpy as np
 import cv2
 import streamlit as st
-from streamlit_webrtc import webrtc_streamer, VideoProcessorBase, WebRtcMode
+from streamlit_webrtc import webrtc_streamer, VideoProcessorBase, WebRtcMode, RTCConfiguration
 from huggingface_hub import hf_hub_download
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.DEBUG)
 
 # Define dataset classes
 classes = ['dengue-regions', 'wet_surface']
@@ -15,46 +19,91 @@ num_classes = len(classes) + 1  # Including background
 # Load the SSD Model
 @st.cache_resource
 def load_model():
-    if 'model' not in st.session_state:
-        model_path = hf_hub_download(repo_id="DhominickJ/MosqScope", filename="mosquito_model.pth")
-        model = ssd300_vgg16(pretrained=True) # Multi-box Algorithm
-        model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
-        # model.load_state_dict(torch.load("./mosquito_model.pth", map_location=torch.device('cpu')))
-        model.eval()
-        st.session_state.model = model
+    model_path = hf_hub_download(repo_id="DhominickJ/MosqScope", filename="mosquito_model.pth")
+    model = ssd300_vgg16(pretrained=False)  # Don't load ImageNet weights
+    in_features = model.roi_heads.box_predictor.cls_score.in_features
+    model.roi_heads.box_predictor = torch.nn.Linear(in_features, num_classes)
+    model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
+    model.eval()
+    return model
 
-model = load_model()
+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
         self.transform = transforms.Compose([
+            transforms.ToPILImage(),
             transforms.Resize((300, 300)),
-            transforms.ToTensor()
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
         ])
 
     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")
-        image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-        image = self.transform(image).unsqueeze(0)
-
-        with torch.no_grad():
-            output = self.model(image)[0]
-
-        # Draw detections
-        for box, label in zip(output["boxes"].cpu().numpy(), output["labels"].cpu().numpy()):
-            x_min, y_min, x_max, y_max = map(int, box)
-            cv2.rectangle(img, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
-            cv2.putText(img, classes[label - 1], (x_min, y_min - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
-
-        return av.VideoFrame.from_ndarray(img, format="bgr24")
-
-# Start WebRTC Streaming
-st.title("Real-Time SSD Object Detection with WebRTC")
-webrtc_streamer(
-    key="ssd-detection",
-    mode=WebRtcMode.SENDRECV,
-    video_processor_factory=SSDVideoProcessor,
+        # Make a copy for drawing
+        display_img = img.copy()
+        
+        # Transform for model
+        image_tensor = self.transform(img).unsqueeze(0)
+
+        try:
+            with torch.no_grad():
+                output = self.model(image_tensor)[0]
+                
+            # 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(output["boxes"], output["labels"], output["scores"]):
+                if score > 0.5:  # Only show confident detections
+                    x_min, y_min, x_max, y_max = map(int, box.cpu().numpy())
+                    # 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.item() - 1]
+                    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")
+
+# 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"]}]},
     media_stream_constraints={"video": True, "audio": False},
 )
+
+# Start WebRTC Streaming with proper error handling
+try:
+    webrtc_ctx = webrtc_streamer(
+        key="ssd-detection",
+        mode=WebRtcMode.SENDRECV,
+        rtc_configuration=rtc_config,
+        video_processor_factory=SSDVideoProcessor,
+        async_processing=True,
+    )
+except Exception as e:
+    st.error(f"WebRTC Error: {e}")
+    st.info("Please try refreshing the page or using a different browser.")