Initial implementation of MosqScope

app.py

@@ -7,27 +7,19 @@ import streamlit as st
 from PIL import Image
 from huggingface_hub import hf_hub_download
 
-# Configure Streamlit UI
-st.title("Mosquito Detection from Camera Capture")
-st.write("Take a picture to detect mosquito breeding sites using SSD.")
-
 # Define dataset classes
 classes = ['dengue-regions', 'wet_surface']
-
-# 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)
-        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)}")
-        return None
-
-model = load_model()
+num_classes = len(classes) + 1  # Including background
+
+# Load Model
+st.title("Real-Time SSD Object Detection")
+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
 
 # Capture Image from Camera
 captured_image = st.camera_input("Take a picture")
@@ -43,28 +35,21 @@ if captured_image is not None:
     ])
     image_tensor = transform(image).unsqueeze(0)
 
-    if model is not None:
-        with torch.no_grad():
-            detections = model(image_tensor)
-
-        boxes = detections[0]['boxes'].cpu().numpy()
-        scores = detections[0]['scores'].cpu().numpy()
-        labels = detections[0]['labels'].cpu().numpy()
-
-        # Convert image to OpenCV format
-        image_np = np.array(image)
-        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
-
-        # Draw detections
-        for box, label, score in zip(boxes, labels, scores):
-            if score > 0.5:  # Confidence threshold
-                x_min, y_min, x_max, y_max = map(int, box)
-                cv2.rectangle(image_np, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
-                label_name = classes[label - 1]
-                cv2.putText(image_np, f"{label_name} {score:.2f}", (x_min, y_min - 5), 
-                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
-
-        # Convert image back to RGB for Streamlit display
-        st.image(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB), caption="Detected Objects", use_column_width=True)
-    else:
-        st.warning("Model not loaded. Unable to process image.")
+    # Convert frame for model
+    image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    image_tensor = transform(image).unsqueeze(0)
+    
+    with torch.no_grad():
+        output = st.session_state.model(image_tensor)[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(frame, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
+        cv2.putText(frame, classes[label - 1], (x_min, y_min - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
+    
+    # Display frame
+    stframe.image(frame, channels="BGR")
+
+cap.release()
+cv2.destroyAllWindows()