Fixing the error with the last instance of onnx

app.py

@@ -1,95 +1,109 @@
 import streamlit as st
-import cv2
 import numpy as np
-import onnxruntime as ort
-from huggingface_hub import hf_hub_download
+import cv2
 from PIL import Image
+from ultralytics import YOLO
+from huggingface_hub import hf_hub_download
 
-# Configurable Variables
-
-MODEL_PATH = hf_hub_download(repo_id="DhominickJ/DengueTect", filename="mosquito_location.onnx")  
-CONF_THRESHOLD = 0.5  # Confidence threshold for detection
-IOU_THRESHOLD = 0.4  # IoU threshold for NMS
-CLASS_NAMES = ["dengue-regions", "wet_surface"]  # Define class names
-IMAGE_SIZE = 640  # YOLO model expects 640x640 images
-
-# Load ONNX model
-session = ort.InferenceSession(MODEL_PATH)
-input_name = session.get_inputs()[0].name
-output_name = session.get_outputs()[0].name
-
-# Function to preprocess image
-def preprocess(image):
-    img = np.array(image)
-    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)  # Convert RGB to BGR for OpenCV
-    img_resized = cv2.resize(img, (IMAGE_SIZE, IMAGE_SIZE))
-    img_transposed = img_resized.transpose(2, 0, 1)  # Change to channel-first
-    img_normalized = img_transposed / 255.0  # Normalize to [0,1]
-    img_input = np.expand_dims(img_normalized.astype(np.float32), axis=0)
-    return img, img_input
 
-# Function to post-process outputs
-def postprocess(outputs, original_shape):
-    outputs = np.squeeze(outputs)  # Remove batch dimension
-    if len(outputs.shape) != 2 or outputs.shape[1] < 6:
-        return [], [], []  # Return empty lists if the output format is incorrect
-    
-    boxes = outputs[:, :4]
-    scores = outputs[:, 4]
-    class_ids = outputs[:, 5].astype(int)
-    
-    # Apply confidence threshold
-    mask = scores >= CONF_THRESHOLD
-    boxes, scores, class_ids = boxes[mask], scores[mask], class_ids[mask]
-    
-    # Ensure valid class IDs
-    valid_mask = (class_ids >= 0) & (class_ids < len(CLASS_NAMES))
-    boxes, scores, class_ids = boxes[valid_mask], scores[valid_mask], class_ids[valid_mask]
-    
-    # Rescale boxes to original image size
-    h, w, _ = original_shape
-    scale_w, scale_h = w / IMAGE_SIZE, h / IMAGE_SIZE
-    boxes[:, [0, 2]] *= scale_w
-    boxes[:, [1, 3]] *= scale_h
-    
-    return boxes.tolist(), scores.tolist(), class_ids.tolist()
+# Define static class names
+CLASS_NAMES = {0: 'dengue-regions', 1: 'wet_surface'}
+CONFIDENCE_VALUE = 0.3
 
-# Function to draw detected bounding boxes
-def draw_boxes(image, boxes, scores, class_ids):
-    for box, score, class_id in zip(boxes, scores, class_ids):
-        x1, y1, x2, y2 = map(int, box)
-        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
-        label = f"{CLASS_NAMES[class_id]}: {score:.2f}"
-        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+# Load YOLOv8 (.pt) model
+@st.cache_resource
+def load_pt_model():
+    # Update the path to your .pt model file
+    model_path = hf_hub_download(repo_id="DhominickJ/DengueTect", filename="best.pt")
+    return YOLO(model_path)
 
-    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert BGR back to RGB before displaying
+# Draw bounding boxes on the image and return the image with detections drawn
+def draw_boxes_pt(image, results):
+    image_cv = np.array(image)
+    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)
+    for result in results:
+        if result.boxes is not None:
+            # Get bounding boxes, confidences, and class IDs
+            boxes = result.boxes.xyxy.cpu().numpy()  # (n,4)
+            confidences = result.boxes.conf.cpu().numpy()  # (n,)
+            classes = result.boxes.cls.cpu().numpy()  # (n,)
+            for box, conf, cls in zip(boxes, confidences, classes):
+                x1, y1, x2, y2 = box.astype(int)
+                label = f"{CLASS_NAMES.get(int(cls), 'Unknown')}: {conf:.2f}"
+                # Set color based on class
+                color = (0, 0, 255) if CLASS_NAMES[int(cls)] == 'dengue-regions' else (255, 0, 0)
+                cv2.rectangle(image_cv, (x1, y1), (x2, y2), color, 3)
+                cv2.putText(image_cv, label, (x1, max(y1-10, 0)), 
+                            cv2.FONT_HERSHEY_SIMPLEX, 1.5, color, 3)
+    # Convert back to RGB
+    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(image_cv)
 
-# Function to perform object detection
-def detect_objects(image):
-    original_img, input_tensor = preprocess(image)
-    outputs = session.run([output_name], {input_name: input_tensor})[0]
-    
-    # Debugging: Show raw outputs
-    st.write("Raw Predictions Shape:", outputs.shape)
-    st.write("Raw Predictions Sample:", outputs.flatten()[:10])
-    
-    boxes, scores, class_ids = postprocess(outputs, original_img.shape)
-    
-    # Show confidence score range
-    if scores:
-        st.write("Min Confidence:", np.min(scores), "Max Confidence:", np.max(scores))
-    else:
-        st.write("No objects detected.")
+# Streamlit App
+st.title("🦟 DengueTect: Object Detection with YOLOv8 (.pt)")
+st.write("Upload an image or capture from your camera to detect mosquito breeding grounds!")
 
-    return draw_boxes(original_img, boxes, scores, class_ids)
+# Choose input method
+option = st.radio("Choose Input Method:", ("Upload Image", "Use Camera"))
 
-# Streamlit UI
-st.title("Mosquito Detection App")
+# Load model once
+model = load_pt_model()
 
-# Capture Image from Camera
-image = st.camera_input("Take a picture")
+if option == "Upload Image":
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file).convert("RGB")
+        st.image(image, caption="Uploaded Image", use_column_width=True)
+        
+        # Run YOLOv8 prediction; the predict() method accepts a NumPy array
+        results = model.predict(source=np.array(image), conf=CONFIDENCE_VALUE, imgsz=640)
+        
+        # Collect detection details (for logging purposes)
+        detections = []
+        for result in results:
+            if result.boxes is not None:
+                boxes = result.boxes.xyxy.cpu().numpy()
+                confidences = result.boxes.conf.cpu().numpy()
+                classes = result.boxes.cls.cpu().numpy()
+                for box, conf, cls in zip(boxes, confidences, classes):
+                    x1, y1, x2, y2 = box.astype(int)
+                    detections.append((int(cls), x1, y1, x2, y2, conf))
+        
+        if detections:
+            st.success(f"✅ Detected {len(detections)} objects!")
+            for det in detections:
+                st.text(f"Detected: {CLASS_NAMES.get(det[0], 'Unknown')}, "
+                        f"BBox: ({det[1]}, {det[2]}), ({det[3]}, {det[4]}), "
+                        f"Confidence: {det[5]:.2f}")
+            result_image = draw_boxes_pt(image, results)
+            st.image(result_image, caption="Detection Results", use_column_width=True)
+        else:
+            st.warning("⚠ No objects detected. Try adjusting confidence threshold.")
 
-if image is not None:
-    image = Image.open(image)
-    detected_img = detect_objects(image)
-    st.image(detected_img, caption="Detected Objects", use_column_width=True)
+elif option == "Use Camera":
+    camera_image = st.camera_input("Capture an image")
+    if camera_image is not None:
+        image = Image.open(camera_image).convert("RGB")
+        st.image(image, caption="Captured Image", use_column_width=True)
+        
+        results = model.predict(source=np.array(image), conf=CONFIDENCE_VALUE, imgsz=640)
+        detections = []
+        for result in results:
+            if result.boxes is not None:
+                boxes = result.boxes.xyxy.cpu().numpy()
+                confidences = result.boxes.conf.cpu().numpy()
+                classes = result.boxes.cls.cpu().numpy()
+                for box, conf, cls in zip(boxes, confidences, classes):
+                    x1, y1, x2, y2 = box.astype(int)
+                    detections.append((int(cls), x1, y1, x2, y2, conf))
+                    
+        if detections:
+            st.success(f"✅ Detected {len(detections)} objects!")
+            for det in detections:
+                st.text(f"Detected: {CLASS_NAMES.get(det[0], 'Unknown')}, "
+                        f"BBox: ({det[1]}, {det[2]}), ({det[3]}, {det[4]}), "
+                        f"Confidence: {det[5]:.2f}")
+            result_image = draw_boxes_pt(image, results)
+            st.image(result_image, caption="Detection Results", use_column_width=True)
+        else:
+            st.warning("⚠ No objects detected. Try adjusting confidence threshold.")

requirements.txt

@@ -1,5 +1,5 @@
 streamlit
 numpy
 opencv-python
-onnxruntime
-Pillow
+Pillow
+ultralytics