Initial implementation of MosqScope

app.py

@@ -7,19 +7,27 @@ 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']
-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
+
+# 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()
 
 # Capture Image from Camera
 captured_image = st.camera_input("Take a picture")
@@ -30,26 +38,33 @@ if captured_image is not None:
 
     # Transform the image for SSD model
     transform = transforms.Compose([
-        transforms.Resize((300, 300)),
+        transforms.Resize((800, 800)),
         transforms.ToTensor()
     ])
     image_tensor = transform(image).unsqueeze(0)
 
-    # 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()
+    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.")