Update app.py

app.py

@@ -19,6 +19,10 @@ The application will display the image with bounding boxes around the objects.
 st.markdown("### Step 1: Upload an Image for Object Detection")
 uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
 
+# Load class labels (COCO dataset)
+with open("coco.names", "r") as f:
+    class_names = [line.strip() for line in f.readlines()]
+
 # Function to detect objects using OpenCV's YOLO
 def detect_objects_with_yolo(image):
     # Load YOLO model configuration and weights
@@ -28,7 +32,7 @@ def detect_objects_with_yolo(image):
 
     # Get the names of the output layers
     layer_names = net.getLayerNames()
-    output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]
+    output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers().flatten()]
 
     # Prepare the image for YOLO
     blob = cv2.dnn.blobFromImage(image, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
@@ -37,8 +41,8 @@ def detect_objects_with_yolo(image):
 
     # Process YOLO outputs
     boxes = []
-    labels = []
     confidences = []
+    class_ids = []
 
     for out in outs:
         for detection in out:
@@ -54,9 +58,21 @@ def detect_objects_with_yolo(image):
                 y = int(center_y - h / 2)
                 boxes.append([x, y, w, h])
                 confidences.append(float(confidence))
-                labels.append(class_id)
+                class_ids.append(class_id)
+
+    # Apply Non-Maximum Suppression (NMS) to remove duplicate boxes
+    indices = cv2.dnn.NMSBoxes(boxes, confidences, score_threshold=0.5, nms_threshold=0.4)
+    final_boxes = []
+    final_labels = []
+    final_confidences = []
 
-    return boxes, labels, confidences
+    if len(indices) > 0:
+        for i in indices.flatten():
+            final_boxes.append(boxes[i])
+            final_labels.append(class_ids[i])
+            final_confidences.append(confidences[i])
+
+    return final_boxes, final_labels, final_confidences
 
 # Display the uploaded image and detect objects
 if uploaded_file is not None:
@@ -70,7 +86,7 @@ if uploaded_file is not None:
     # Draw bounding boxes on the image
     for i, box in enumerate(boxes):
         x, y, w, h = box
-        label = labels[i]
+        label = class_names[labels[i]] if labels[i] < len(class_names) else f"ID {labels[i]}"
         confidence = confidences[i]
         cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)
         cv2.putText(image, f"{label} {confidence:.2f}", (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
@@ -81,11 +97,11 @@ if uploaded_file is not None:
 
     # Step 3: Display the count of detected objects dynamically
     st.markdown("### Step 3: Detected Objects Count")
-    label_counts = Counter(labels)
+    label_counts = Counter([class_names[l] for l in labels])
 
     # Display counts in a well-formatted table
     for obj, count in label_counts.items():
-        st.markdown(f"**Object ID {obj}**: {count}")
+        st.markdown(f"**{obj}**: {count}")
 
     st.markdown("""
     ---
@@ -99,4 +115,4 @@ if uploaded_file is not None:
 st.markdown("""
 ---
 Made with ❤️ by [SenasuDemir](https://github.com/SenasuDemir).
-""")
+""")