Update app.py

added overlapping threshold slider

app.py

@@ -17,7 +17,7 @@ def save_uploaded_file(uploaded_file):
         tmp_file.write(uploaded_file.getbuffer())
         return tmp_file.name
 
-def apply_confidence_threshold(result, conf_threshold):
+def apply_confidence_threshold(result, conf_threshold, iou_threshold=0.45):
     """Apply confidence threshold by modifying the result's boxes directly."""
     try:
         # If there are no boxes, or the boxes have no confidence values, just return the original image
@@ -38,14 +38,19 @@ def apply_confidence_threshold(result, conf_threshold):
             img_with_boxes = result.orig_img.copy()
         else:
             # Fallback to plot method if orig_img is not available
-            return Image.fromarray(np.array(result.plot(conf=conf_threshold))), valid_detections
+            return Image.fromarray(np.array(result.plot(conf=conf_threshold, iou=iou_threshold))), valid_detections
         
         # Only proceed with drawing if there are valid detections
         if valid_detections > 0:
             # Create mask of boxes to keep
             mask = confs >= conf_threshold
             
-            # For each valid box, draw it on the image
+            # Apply non-maximum suppression if we have xyxy boxes
+            boxes_to_draw = []
+            class_ids = []
+            confidences = []
+            
+            # Collect all valid boxes
             for i, is_valid in enumerate(mask):
                 if is_valid:
                     try:
@@ -71,7 +76,7 @@ def apply_confidence_threshold(result, conf_threshold):
                             # If we can't get box coordinates, skip this box
                             continue
                         
-                        # Get class ID and name
+                        # Get class ID
                         if hasattr(result.boxes, "cls"):
                             if hasattr(result.boxes.cls, "cpu"):
                                 cls_id = int(result.boxes.cls[i].cpu().item())
@@ -83,12 +88,59 @@ def apply_confidence_threshold(result, conf_threshold):
                         # Get confidence
                         conf = confs[i]
                         
-                        # Get class name
-                        if hasattr(result, 'names') and result.names and cls_id in result.names:
-                            cls_name = result.names[cls_id]
-                        else:
-                            cls_name = f"class_{cls_id}"
+                        # Add to our collection
+                        boxes_to_draw.append(box)
+                        class_ids.append(cls_id)
+                        confidences.append(conf)
                         
+                    except Exception as e:
+                        # If any error occurs for a specific box, just skip it
+                        st.error(f"Error processing a detection box: {str(e)}")
+                        continue
+            
+            # Apply non-maximum suppression if we have collected boxes
+            if boxes_to_draw:
+                try:
+                    # Convert to numpy arrays
+                    boxes_np = np.array(boxes_to_draw)
+                    class_ids_np = np.array(class_ids)
+                    confidences_np = np.array(confidences)
+                    
+                    # Apply NMS by class
+                    unique_classes = np.unique(class_ids_np)
+                    final_boxes = []
+                    final_classes = []
+                    final_confs = []
+                    
+                    for cls in unique_classes:
+                        # Get indices for this class
+                        indices = np.where(class_ids_np == cls)[0]
+                        if len(indices) == 0:
+                            continue
+                            
+                        class_boxes = boxes_np[indices]
+                        class_confs = confidences_np[indices]
+                        
+                        # Apply NMS
+                        keep_indices = cv2.dnn.NMSBoxes(
+                            class_boxes.tolist(), 
+                            class_confs.tolist(), 
+                            score_threshold=conf_threshold, 
+                            nms_threshold=iou_threshold
+                        )
+                        
+                        # Add kept boxes to final lists
+                        if len(keep_indices) > 0:
+                            if isinstance(keep_indices[0], np.ndarray):  # Handle different return formats
+                                keep_indices = keep_indices.flatten()
+                                
+                            for idx in keep_indices:
+                                final_boxes.append(class_boxes[idx])
+                                final_classes.append(cls)
+                                final_confs.append(class_confs[idx])
+                    
+                    # Now draw only the final boxes
+                    for box, cls_id, conf in zip(final_boxes, final_classes, final_confs):
                         # Make sure box coordinates are within image bounds
                         h, w = img_with_boxes.shape[:2]
                         box[0] = max(0, min(box[0], w-1))
@@ -96,6 +148,12 @@ def apply_confidence_threshold(result, conf_threshold):
                         box[2] = max(0, min(box[2], w-1))
                         box[3] = max(0, min(box[3], h-1))
                         
+                        # Get class name
+                        if hasattr(result, 'names') and result.names and cls_id in result.names:
+                            cls_name = result.names[cls_id]
+                        else:
+                            cls_name = f"class_{cls_id}"
+                            
                         # Draw the box
                         color = (0, 255, 0)  # Green box
                         cv2.rectangle(img_with_boxes, (box[0], box[1]), (box[2], box[3]), color, 2)
@@ -109,10 +167,15 @@ def apply_confidence_threshold(result, conf_threshold):
                         text_x = box[0]
                         text_y = max(box[1] - 10, text_size[1])
                         cv2.putText(img_with_boxes, label, (text_x, text_y), font, 0.6, color, 2)
-                    except Exception as e:
-                        # If any error occurs for a specific box, just skip it
-                        st.error(f"Error processing a detection box: {str(e)}")
-                        continue
+                        
+                    # Update valid_detections to reflect NMS results
+                    valid_detections = len(final_boxes)
+                    
+                except Exception as nms_error:
+                    # If NMS fails, fall back to original drawing code
+                    st.warning(f"NMS processing failed, falling back to simple filtering: {str(nms_error)}")
+                    # The original boxes will be drawn in the fallback code
+                    pass
         
         # Convert back to PIL Image for streamlit display
         img_pil = Image.fromarray(img_with_boxes)
@@ -122,7 +185,7 @@ def apply_confidence_threshold(result, conf_threshold):
         # If anything fails in the custom drawing, fall back to the model's built-in plot method
         try:
             # Try using the built-in plot method with the threshold
-            annotated_img = result.plot(conf=conf_threshold)
+            annotated_img = result.plot(conf=conf_threshold, iou=iou_threshold)
             if isinstance(annotated_img, np.ndarray):
                 img_pil = Image.fromarray(annotated_img)
             else:
@@ -284,17 +347,29 @@ def yolo_inference_tool():
             key="single_model_conf_threshold"
         )
         
-        # Display annotated images using the current threshold
+        # Add IoU threshold slider for NMS
+        st.subheader("Overlapping (IoU) Threshold")
+        iou_threshold = st.slider(
+            "Adjust IoU threshold for non-maximum suppression", 
+            min_value=0.1, 
+            max_value=1.0, 
+            value=0.45,  # Default NMS value
+            step=0.05,
+            key="single_model_iou_threshold",
+            help="Higher values allow more overlapping boxes. Lower values keep only the most confident box in overlapping groups."
+        )
+        
+        # Display annotated images using the current thresholds
         st.subheader("Annotated Images")
         for img_name, r in st.session_state.single_model_results.items():
             try:
-                # Apply confidence threshold and get processed image
-                processed_img, valid_detections = apply_confidence_threshold(r, conf_threshold)
+                # Apply confidence and IoU thresholds and get processed image
+                processed_img, valid_detections = apply_confidence_threshold(r, conf_threshold, iou_threshold)
                 
                 # Display the image
                 st.image(
                     processed_img, 
-                    caption=f"{img_name} (Threshold: {conf_threshold:.2f}, Detections: {valid_detections})", 
+                    caption=f"{img_name} (Conf: {conf_threshold:.2f}, IoU: {iou_threshold:.2f}, Detections: {valid_detections})", 
                     use_container_width=True
                 )
             except Exception as e:
@@ -520,6 +595,18 @@ def yolo_model_comparison_tool():
             step=0.05,
             key="multi_model_conf_threshold"
         )
+        
+        # Add IoU threshold slider for NMS
+        st.subheader("Overlapping (IoU) Threshold")
+        comp_iou_threshold = st.slider(
+            "Adjust IoU threshold for non-maximum suppression across all models", 
+            min_value=0.1, 
+            max_value=1.0, 
+            value=0.45,  # Default NMS value
+            step=0.05,
+            key="multi_model_iou_threshold",
+            help="Higher values allow more overlapping boxes. Lower values keep only the most confident box in overlapping groups."
+        )
 
         # Display annotated images in a grid (row = image, column = model)
         st.subheader("Annotated Images Grid (Row = Image, Column = Model)")
@@ -540,12 +627,12 @@ def yolo_model_comparison_tool():
                     continue
 
                 try:
-                    # Apply confidence threshold and get processed image
-                    processed_img, valid_detections = apply_confidence_threshold(r, comp_conf_threshold)
+                    # Apply confidence and IoU thresholds and get processed image
+                    processed_img, valid_detections = apply_confidence_threshold(r, comp_conf_threshold, comp_iou_threshold)
                     
                     col.image(
                         processed_img,
-                        caption=f"{model_name} (Threshold: {comp_conf_threshold:.2f}, Detections: {valid_detections})",
+                        caption=f"{model_name} (Conf: {comp_conf_threshold:.2f}, IoU: {comp_iou_threshold:.2f}, Det: {valid_detections})",
                         use_container_width=True
                     )
                 except Exception as e: