Update app.py

added threshold slider

app.py

@@ -7,6 +7,8 @@ from PIL import Image
 from ultralytics import YOLO
 import requests
 from io import BytesIO
+import copy
+import cv2
 
 
 def save_uploaded_file(uploaded_file):
@@ -15,6 +17,138 @@ def save_uploaded_file(uploaded_file):
         tmp_file.write(uploaded_file.getbuffer())
         return tmp_file.name
 
+def apply_confidence_threshold(result, conf_threshold):
+    """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
+        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)
+        
+        # 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))), 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
+            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 and name
+                        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]
+                        
+                        # 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}"
+                        
+                        # 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))
+                        
+                        # 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)
+                    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
+        
+        # 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
+        try:
+            # Try using the built-in plot method with the threshold
+            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
+            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()
+                else:
+                    confs = result.boxes.conf
+                valid_detections = sum(confs >= conf_threshold)
+            else:
+                valid_detections = 0
+                
+            return img_pil, valid_detections
+        except Exception as nested_e:
+            # If even the fallback fails, return the original image without annotations
+            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
+            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)
+            return Image.fromarray(blank_img), 0
+
 def yolo_inference_tool():
     st.header("YOLO Model Inference Tool")
     st.write(
@@ -22,6 +156,12 @@ def yolo_inference_tool():
         "You can either upload images or provide an image URL."
     )
     
+    # Initialize session state for storing inference results
+    if 'single_model_results' not in st.session_state:
+        st.session_state.single_model_results = None
+    if 'single_model_metrics' not in st.session_state:
+        st.session_state.single_model_metrics = None
+    
     # Allow multiple images upload
     uploaded_files = st.file_uploader(
         "Upload Images", type=["jpg", "jpeg", "png"], key="inference_images", accept_multiple_files=True
@@ -89,7 +229,8 @@ def yolo_inference_tool():
                 continue
 
             try:
-                result = model(np.array(pil_img))
+                # Run inference with the lowest possible confidence to capture all detections
+                result = model(np.array(pil_img), conf=0.01)
             except Exception as e:
                 st.error(f"Inference error on image {getattr(img_file, 'name', 'Unknown')}: {e}")
                 continue
@@ -100,10 +241,14 @@ def yolo_inference_tool():
             # Get inference time from r.speed, if available
             inference_time = r.speed.get('inference', None) if isinstance(r.speed, dict) else None
             # Compute detection count and average confidence if detections exist
-            if r.boxes is not None and len(r.boxes) > 0:
+            if hasattr(r, 'boxes') and r.boxes is not None and len(r.boxes) > 0:
                 detection_count = len(r.boxes)
-                confs = r.boxes.conf.cpu().numpy() if hasattr(r.boxes.conf, "cpu") else r.boxes.conf
-                avg_conf = float(np.mean(confs))
+                if hasattr(r.boxes.conf, "cpu"):
+                    confs = r.boxes.conf.cpu().numpy()
+                    avg_conf = float(np.mean(confs))
+                else:
+                    confs = r.boxes.conf
+                    avg_conf = float(np.mean(confs))
             else:
                 detection_count = 0
                 avg_conf = 0.0
@@ -117,20 +262,44 @@ def yolo_inference_tool():
         
         eta_placeholder.empty()
         
-        # Display per-image metrics if collected
+        # Store results in session state for persistence
+        st.session_state.single_model_results = image_results
+        st.session_state.single_model_metrics = metrics
+
+    # Display results if available in session state (either from button click or slider change)
+    if st.session_state.single_model_metrics is not None:
+        # Display per-image metrics
         st.subheader("Inference Metrics")
-        if metrics:
-            df_metrics = pd.DataFrame(metrics)
-            st.dataframe(df_metrics, use_container_width=True)
+        df_metrics = pd.DataFrame(st.session_state.single_model_metrics)
+        st.dataframe(df_metrics, use_container_width=True)
+        
+        # Add a confidence threshold slider
+        st.subheader("Confidence Threshold")
+        conf_threshold = st.slider(
+            "Adjust confidence threshold", 
+            min_value=0.0, 
+            max_value=1.0, 
+            value=0.25,  # Default value
+            step=0.05,
+            key="single_model_conf_threshold"
+        )
         
-        # Display annotated images using pil=True (ensuring RGB)
+        # Display annotated images using the current threshold
         st.subheader("Annotated Images")
-        for img_name, r in image_results.items():
+        for img_name, r in st.session_state.single_model_results.items():
             try:
-                annotated_img = r.plot(conf=True, boxes=True, labels=True, pil=True)
-                st.image(annotated_img, caption=img_name, use_container_width=True)
+                # Apply confidence threshold and get processed image
+                processed_img, valid_detections = apply_confidence_threshold(r, conf_threshold)
+                
+                # Display the image
+                st.image(
+                    processed_img, 
+                    caption=f"{img_name} (Threshold: {conf_threshold:.2f}, Detections: {valid_detections})", 
+                    use_container_width=True
+                )
             except Exception as e:
                 st.error(f"Error generating annotated image for {img_name}: {e}")
+                st.error(str(e))
 
 
 def yolo_model_comparison_tool():
@@ -147,6 +316,16 @@ def yolo_model_comparison_tool():
         "and compute a Weighted Score that balances these factors.\n\n"
         "A progress bar and ETA are shown in real time after you click Submit."
     )
+    
+    # Initialize session state for storing model comparison results
+    if 'model_agg_data' not in st.session_state:
+        st.session_state.model_agg_data = None
+    if 'model_image_results' not in st.session_state:
+        st.session_state.model_image_results = None
+    if 'model_metrics_df' not in st.session_state:
+        st.session_state.model_metrics_df = None
+    if 'best_model_info' not in st.session_state:
+        st.session_state.best_model_info = None
 
     images = st.file_uploader("Upload Images", type=["jpg", "jpeg", "png"], key="comparison_images", accept_multiple_files=True)
     model_files = st.file_uploader("Upload YOLO models (.pt)", type=["pt"], key="comparison_models", accept_multiple_files=True)
@@ -212,7 +391,8 @@ def yolo_model_comparison_tool():
 
                 # Run inference
                 try:
-                    result = model(np_img)
+                    # Use low confidence to capture all detections
+                    result = model(np_img, conf=0.01)
                 except Exception as e:
                     st.error(f"Inference error for model {model_file.name} on {img_file.name}: {e}")
                     continue
@@ -225,11 +405,14 @@ def yolo_model_comparison_tool():
                     total_inference_time += r.speed["inference"]
 
                 # Count detections & confidence
-                if r.boxes is not None:
+                if hasattr(r, 'boxes') and r.boxes is not None and len(r.boxes) > 0:
                     det_count = len(r.boxes)
                     total_detections += det_count
                     if det_count > 0:
-                        confs = r.boxes.conf.cpu().numpy() if hasattr(r.boxes.conf, "cpu") else r.boxes.conf
+                        if hasattr(r.boxes.conf, "cpu"):
+                            confs = r.boxes.conf.cpu().numpy()
+                        else:
+                            confs = r.boxes.conf
                         sum_confidences += confs.sum()
                         total_conf_count += det_count
 
@@ -254,9 +437,7 @@ def yolo_model_comparison_tool():
 
         # Display aggregated metrics
         df = pd.DataFrame(model_agg_data.values())
-        st.subheader("Aggregated Metrics (Across All Images)")
-        st.dataframe(df, use_container_width=True)
-
+        
         # Weighted Scoring with reciprocal-based speed
         detection_max = df["Total Detections"].max()
         confidence_max = df["Average Confidence"].max()
@@ -281,6 +462,25 @@ def yolo_model_comparison_tool():
             gamma_speed     * df["Speed Norm"]
         )
 
+        # Identify best overall model (highest Weighted Score)
+        best_idx = df["Weighted Score"].idxmax()
+        best_model = df.loc[best_idx, "Model File"]
+        best_score = df.loc[best_idx, "Weighted Score"]
+        
+        # Store results in session state
+        st.session_state.model_agg_data = model_agg_data
+        st.session_state.model_image_results = model_image_results
+        st.session_state.model_metrics_df = df
+        st.session_state.best_model_info = (best_model, best_score)
+
+    # Display results if available in session state
+    if st.session_state.model_metrics_df is not None:
+        df = st.session_state.model_metrics_df
+        best_model, best_score = st.session_state.best_model_info
+        
+        st.subheader("Aggregated Metrics (Across All Images)")
+        st.dataframe(df, use_container_width=True)
+
         st.subheader("Weighted Score Analysis")
         st.write(f"Weights: Detection={alpha_detection}, Confidence={beta_confidence}, Speed={gamma_speed}")
         st.dataframe(df[[
@@ -295,11 +495,6 @@ def yolo_model_comparison_tool():
             "Weighted Score"
         ]], use_container_width=True)
 
-        # Identify best overall model (highest Weighted Score)
-        best_idx = df["Weighted Score"].idxmax()
-        best_model = df.loc[best_idx, "Model File"]
-        best_score = df.loc[best_idx, "Weighted Score"]
-
         st.markdown(f"""
         **Best Overall Model** based on Weighted Score:
         **{best_model}** (Score: {best_score:.3f}).
@@ -315,29 +510,47 @@ def yolo_model_comparison_tool():
         - Increase **Speed** weight if you need real‐time inference.
         """)
 
+        # Add a confidence threshold slider
+        st.subheader("Confidence Threshold")
+        comp_conf_threshold = st.slider(
+            "Adjust confidence threshold for all models", 
+            min_value=0.0, 
+            max_value=1.0, 
+            value=0.25,  # Default value
+            step=0.05,
+            key="multi_model_conf_threshold"
+        )
+
         # Display annotated images in a grid (row = image, column = model)
         st.subheader("Annotated Images Grid (Row = Image, Column = Model)")
-        model_names_sorted = sorted(model_agg_data.keys())
-
-        for img_file in images:
-            st.markdown(f"### Image: {img_file.name}")
+        model_names_sorted = sorted(st.session_state.model_agg_data.keys())
+        
+        # Extract the image file names from the stored results
+        image_names = set()
+        for model_results in st.session_state.model_image_results.values():
+            image_names.update(model_results.keys())
+        
+        for img_name in sorted(image_names):
+            st.markdown(f"### Image: {img_name}")
             columns = st.columns(len(model_names_sorted))
             for col, model_name in zip(columns, model_names_sorted):
-                r = model_image_results.get(model_name, {}).get(img_file.name, None)
+                r = st.session_state.model_image_results.get(model_name, {}).get(img_name, None)
                 if r is None:
                     col.write(f"No results for {model_name}")
                     continue
 
                 try:
-                    # Return a PIL image in correct RGB color space
-                    annotated_img_pil = r.plot(conf=True, boxes=True, labels=True, pil=True)
+                    # Apply confidence threshold and get processed image
+                    processed_img, valid_detections = apply_confidence_threshold(r, comp_conf_threshold)
+                    
                     col.image(
-                        annotated_img_pil,
-                        caption=f"{model_name}",
+                        processed_img,
+                        caption=f"{model_name} (Threshold: {comp_conf_threshold:.2f}, Detections: {valid_detections})",
                         use_container_width=True
                     )
                 except Exception as e:
                     col.error(f"Error annotating image for {model_name}: {e}")
+                    col.error(str(e))
 
 def main():
     st.sidebar.title("Navigation")