Update app.py

fixed the overlapping threshold slider.

app.py

@@ -23,175 +23,228 @@ def apply_confidence_threshold(result, conf_threshold, iou_threshold=0.45):
         # If there are no boxes, or the boxes have no confidence values, just return the original image
         if not hasattr(result, 'boxes') or result.boxes is None or len(result.boxes) == 0:
             return Image.fromarray(result.orig_img), 0
-        
+            
         # Get the confidence values
         if hasattr(result.boxes.conf, "cpu"):
             confs = result.boxes.conf.cpu().numpy()
         else:
             confs = result.boxes.conf
-        
-        # Count valid detections for display purposes
-        valid_detections = sum(confs >= conf_threshold)
+            
+        # First filter by confidence threshold
+        conf_mask = confs >= conf_threshold
         
         # Create a completely new plot with only the boxes that meet the threshold
         if hasattr(result, 'orig_img'):
             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, iou=iou_threshold))), valid_detections
+            try:
+                # First try the combined approach
+                return Image.fromarray(np.array(result.plot(conf=conf_threshold, iou=iou_threshold))), sum(conf_mask)
+            except:
+                # Fallback to just confidence if iou param is not supported
+                return Image.fromarray(np.array(result.plot(conf=conf_threshold))), sum(conf_mask)
+        
+        # Collect all boxes that meet confidence threshold
+        filtered_boxes = []
+        filtered_classes = []
+        filtered_confs = []
+        
+        for i in range(len(confs)):
+            if confs[i] < conf_threshold:
+                continue
+                
+            try:
+                # Get the box coordinates (handle different formats)
+                if hasattr(result.boxes, "xyxy"):
+                    if hasattr(result.boxes.xyxy, "cpu"):
+                        box = result.boxes.xyxy[i].cpu().numpy().astype(float)
+                    else:
+                        box = result.boxes.xyxy[i].astype(float)
+                elif hasattr(result.boxes, "xywh"):
+                    if hasattr(result.boxes.xywh, "cpu"):
+                        xywh = result.boxes.xywh[i].cpu().numpy().astype(float)
+                    else:
+                        xywh = result.boxes.xywh[i].astype(float)
+                    box = np.array([
+                        xywh[0] - xywh[2]/2,  # x1 = x - w/2
+                        xywh[1] - xywh[3]/2,  # y1 = y - h/2
+                        xywh[0] + xywh[2]/2,  # x2 = x + w/2
+                        xywh[1] + xywh[3]/2   # y2 = y + h/2
+                    ]).astype(float)
+                else:
+                    continue  # Skip if no box format available
+                
+                # Get class ID
+                if hasattr(result.boxes, "cls"):
+                    if hasattr(result.boxes.cls, "cpu"):
+                        cls_id = int(result.boxes.cls[i].cpu().item())
+                    else:
+                        cls_id = int(result.boxes.cls[i])
+                else:
+                    cls_id = 0  # Default class ID if not available
+                
+                # Store the box, class, and confidence
+                filtered_boxes.append(box)
+                filtered_classes.append(cls_id)
+                filtered_confs.append(confs[i])
+                
+            except Exception as e:
+                st.error(f"Error processing detection box: {str(e)}")
+                continue
+        
+        if not filtered_boxes:
+            # No boxes passed the confidence threshold
+            return Image.fromarray(img_with_boxes), 0
+        
+        # Convert to numpy arrays for processing
+        boxes_array = np.array(filtered_boxes)
+        classes_array = np.array(filtered_classes)
+        confs_array = np.array(filtered_confs)
         
-        # Only proceed with drawing if there are valid detections
-        if valid_detections > 0:
-            # Create mask of boxes to keep
-            mask = confs >= conf_threshold
+        # Get unique classes for per-class NMS
+        unique_classes = np.unique(classes_array)
+        
+        # Final boxes to draw after NMS
+        final_boxes = []
+        final_classes = []
+        final_confs = []
+        
+        # Helper function to calculate IoU between two boxes
+        def calculate_iou(box1, box2):
+            # Calculate intersection area
+            x1 = max(box1[0], box2[0])
+            y1 = max(box1[1], box2[1])
+            x2 = min(box1[2], box2[2])
+            y2 = min(box1[3], box2[3])
             
-            # Apply non-maximum suppression if we have xyxy boxes
-            boxes_to_draw = []
-            class_ids = []
-            confidences = []
+            if x2 < x1 or y2 < y1:
+                return 0.0  # No intersection
             
-            # Collect all valid boxes
-            for i, is_valid in enumerate(mask):
-                if is_valid:
-                    try:
-                        # Get the box coordinates (handle different formats)
-                        if hasattr(result.boxes, "xyxy"):
-                            if hasattr(result.boxes.xyxy, "cpu"):
-                                box = result.boxes.xyxy[i].cpu().numpy().astype(int)
-                            else:
-                                box = result.boxes.xyxy[i].astype(int)
-                        elif hasattr(result.boxes, "xywh"):  # Handle xywh format if that's what's available
-                            if hasattr(result.boxes.xywh, "cpu"):
-                                xywh = result.boxes.xywh[i].cpu().numpy().astype(int)
-                            else:
-                                xywh = result.boxes.xywh[i].astype(int)
-                            # Convert xywh to xyxy: [x, y, w, h] -> [x1, y1, x2, y2]
-                            box = np.array([
-                                xywh[0] - xywh[2]//2,  # x1 = x - w/2
-                                xywh[1] - xywh[3]//2,  # y1 = y - h/2
-                                xywh[0] + xywh[2]//2,  # x2 = x + w/2
-                                xywh[1] + xywh[3]//2   # y2 = y + h/2
-                            ]).astype(int)
-                        else:
-                            # If we can't get box coordinates, skip this box
-                            continue
-                        
-                        # Get class ID
-                        if hasattr(result.boxes, "cls"):
-                            if hasattr(result.boxes.cls, "cpu"):
-                                cls_id = int(result.boxes.cls[i].cpu().item())
-                            else:
-                                cls_id = int(result.boxes.cls[i])
-                        else:
-                            cls_id = 0  # Default class ID if not available
-                        
-                        # Get confidence
-                        conf = confs[i]
-                        
-                        # 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
+            intersection_area = (x2 - x1) * (y2 - y1)
             
-            # 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))
-                        box[1] = max(0, min(box[1], h-1))
-                        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)
-                        
-                        # Add label with confidence
-                        label = f"{cls_name} {conf:.2f}"
-                        font = cv2.FONT_HERSHEY_SIMPLEX
-                        # Calculate text size to place it properly
-                        text_size = cv2.getTextSize(label, font, 0.6, 2)[0]
-                        # Ensure label is drawn within image bounds
-                        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)
-                        
-                    # 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
+            # Calculate union area
+            box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
+            box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
+            union_area = box1_area + box2_area - intersection_area
+            
+            # Return IoU
+            if union_area <= 0:
+                return 0.0
+            return intersection_area / union_area
+        
+        # Apply NMS per class as shown in the diagram
+        for cls in unique_classes:
+            # Get all boxes for this class
+            class_indices = np.where(classes_array == cls)[0]
+            if len(class_indices) == 0:
+                continue
+                
+            # Get boxes and scores for this class
+            class_boxes = boxes_array[class_indices]
+            class_scores = confs_array[class_indices]
+            
+            # We'll keep track of which boxes to keep
+            keep_boxes = []
+            
+            # While we still have boxes to process
+            while len(class_indices) > 0:
+                # Find the box with highest confidence
+                max_conf_idx = np.argmax(class_scores)
+                max_conf_box = class_boxes[max_conf_idx]
+                max_conf = class_scores[max_conf_idx]
+                
+                # Add this box to our final list
+                keep_boxes.append(class_indices[max_conf_idx])
+                
+                # Remove this box from consideration
+                class_boxes = np.delete(class_boxes, max_conf_idx, axis=0)
+                class_scores = np.delete(class_scores, max_conf_idx)
+                class_indices = np.delete(class_indices, max_conf_idx)
+                
+                # If no boxes left, we're done with this class
+                if len(class_indices) == 0:
+                    break
+                
+                # Calculate IoU of the saved box with the rest
+                ious = np.array([calculate_iou(max_conf_box, box) for box in class_boxes])
+                
+                # Remove boxes with IoU > threshold
+                boxes_to_keep = ious <= iou_threshold
+                class_boxes = class_boxes[boxes_to_keep]
+                class_scores = class_scores[boxes_to_keep]
+                class_indices = class_indices[boxes_to_keep]
+            
+            # Add all kept boxes for this class to our final lists
+            for idx in keep_boxes:
+                final_boxes.append(filtered_boxes[idx])
+                final_classes.append(filtered_classes[idx])
+                final_confs.append(filtered_confs[idx])
+        
+        # Count valid detections after NMS
+        valid_detections = len(final_boxes)
+        
+        # Draw all final boxes
+        for i, (box, cls_id, conf) in enumerate(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))
+            box[1] = max(0, min(box[1], h-1))
+            box[2] = max(0, min(box[2], w-1))
+            box[3] = max(0, min(box[3], h-1))
+            
+            # Convert to integers for drawing
+            box = box.astype(int)
+            
+            # 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}"
+            
+            # Create a deterministic color based on class ID
+            # Fixed color per class for consistency
+            color_r = (cls_id * 100 + 50) % 255
+            color_g = (cls_id * 50 + 170) % 255
+            color_b = (cls_id * 80 + 90) % 255
+            color = (color_b, color_g, color_r)  # BGR format for OpenCV
+            
+            # Draw rectangle
+            cv2.rectangle(img_with_boxes, (box[0], box[1]), (box[2], box[3]), color, 2)
+            
+            # Add label with confidence
+            label = f"{cls_name} {conf:.2f}"
+            font = cv2.FONT_HERSHEY_SIMPLEX
+            text_size = cv2.getTextSize(label, font, 0.5, 2)[0]
+            
+            # Create filled rectangle for text background
+            rect_y1 = max(0, box[1] - text_size[1] - 10)
+            cv2.rectangle(img_with_boxes, (box[0], rect_y1), 
+                          (box[0] + text_size[0], box[1]), color, -1)
+            
+            # Draw text with white color
+            cv2.putText(img_with_boxes, label, (box[0], box[1] - 5), 
+                        font, 0.5, (255, 255, 255), 1)
+        
+        # Return the annotated image and detection count
+        return Image.fromarray(img_with_boxes), valid_detections
         
-        # Convert back to PIL Image for streamlit display
-        img_pil = Image.fromarray(img_with_boxes)
-        return img_pil, valid_detections
-    
     except Exception as e:
-        # If anything fails in the custom drawing, fall back to the model's built-in plot method
+        # If our custom implementation fails, try using 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, iou=iou_threshold)
+            try:
+                # Try with both parameters if supported
+                annotated_img = result.plot(conf=conf_threshold, iou=iou_threshold)
+            except:
+                # Fallback to just confidence parameter
+                annotated_img = result.plot(conf=conf_threshold)
+                
             if isinstance(annotated_img, np.ndarray):
                 img_pil = Image.fromarray(annotated_img)
             else:
                 img_pil = annotated_img
             
-            # Count detections meeting threshold
+            # Count detections meeting the confidence threshold
             if hasattr(result, 'boxes') and result.boxes is not None and len(result.boxes) > 0:
                 if hasattr(result.boxes.conf, "cpu"):
                     confs = result.boxes.conf.cpu().numpy()
@@ -202,11 +255,13 @@ def apply_confidence_threshold(result, conf_threshold, iou_threshold=0.45):
                 valid_detections = 0
                 
             return img_pil, valid_detections
+            
         except Exception as nested_e:
-            # If even the fallback fails, return the original image without annotations
+            # Last resort: return the original image
             if hasattr(result, 'orig_img'):
                 return Image.fromarray(result.orig_img), 0
-            # If we can't even get the original image, create a blank one with error text
+                
+            # If even that fails, create a blank image with error message
             blank_img = np.zeros((400, 600, 3), dtype=np.uint8)
             cv2.putText(blank_img, f"Error: {str(e)}", (20, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
             cv2.putText(blank_img, "Could not render annotations", (20, 100), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
@@ -351,12 +406,12 @@ def yolo_inference_tool():
         st.subheader("Overlapping (IoU) Threshold")
         iou_threshold = st.slider(
             "Adjust IoU threshold for non-maximum suppression", 
-            min_value=0.1, 
+            min_value=0.0, 
             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."
+            help="Controls how overlapping boxes are filtered. Lower values (0.1-0.3) remove more overlapping boxes, higher values (0.7-0.9) allow more overlaps. The standard YOLO default is 0.45."
         )
         
         # Display annotated images using the current thresholds
@@ -600,12 +655,12 @@ def yolo_model_comparison_tool():
         st.subheader("Overlapping (IoU) Threshold")
         comp_iou_threshold = st.slider(
             "Adjust IoU threshold for non-maximum suppression across all models", 
-            min_value=0.1, 
+            min_value=0.0, 
             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."
+            help="Controls how overlapping boxes are filtered. Lower values (0.1-0.3) remove more overlapping boxes, higher values (0.7-0.9) allow more overlaps. The standard YOLO default is 0.45."
         )
 
         # Display annotated images in a grid (row = image, column = model)