import streamlit as st import tempfile import pandas as pd import numpy as np import time from PIL import Image from ultralytics import YOLO import requests from io import BytesIO import copy import cv2 def save_uploaded_file(uploaded_file): """Save an uploaded file to a temporary file and return its path.""" with tempfile.NamedTemporaryFile(delete=False, suffix=uploaded_file.name) as tmp_file: tmp_file.write(uploaded_file.getbuffer()) return tmp_file.name 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 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 # 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 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) # 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]) if x2 < x1 or y2 < y1: return 0.0 # No intersection intersection_area = (x2 - x1) * (y2 - y1) # 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 except Exception as e: # If our custom implementation fails, try using the model's built-in plot method try: 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 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() 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: # Last resort: return the original image if hasattr(result, 'orig_img'): return Image.fromarray(result.orig_img), 0 # 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) return Image.fromarray(blank_img), 0 def yolo_inference_tool(): st.header("YOLO Model Inference Tool") st.write( "Upload one or more images and a YOLO model (.pt) file to run inference and view detailed results. " "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 ) # Text input for a single image URL (you could expand this to multiple URLs if needed) url_input = st.text_input("Enter image URL (optional)", key="inference_url") # Combine uploaded files and URL image into a single list images = [] if uploaded_files: images.extend(uploaded_files) if url_input and url_input.strip(): try: response = requests.get(url_input) if response.status_code == 200: image_bytes = BytesIO(response.content) # Assign a name attribute for consistency image_bytes.name = url_input images.append(image_bytes) else: st.error("Failed to fetch image from URL.") except Exception as e: st.error(f"Error fetching image from URL: {e}") model_file = st.file_uploader("Upload YOLO model (.pt)", type=["pt"], key="inference_model") if st.button("Submit (Single-Model Inference)"): if not images or not model_file: st.error("Please upload at least one image (or provide an image URL) and a model.") return # Save and load the model file model_path = save_uploaded_file(model_file) try: model = YOLO(model_path) except Exception as e: st.error(f"Error loading model: {e}") return total_images = len(images) progress_bar = st.progress(0) eta_placeholder = st.empty() start_time = time.time() steps_done = 0 # Dictionaries to store inference results and metrics image_results = {} metrics = [] for img_file in images: steps_done += 1 fraction_done = steps_done / total_images progress_bar.progress(fraction_done) elapsed_time = time.time() - start_time time_per_step = elapsed_time / steps_done remaining_steps = total_images - steps_done eta_seconds = remaining_steps * time_per_step eta_placeholder.info(f"Progress: {fraction_done:.1%}. ETA: ~{eta_seconds:.1f} s") try: pil_img = Image.open(img_file).convert("RGB") except Exception as e: st.error(f"Error reading image {getattr(img_file, 'name', 'Unknown')}: {e}") continue try: # 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 r = result[0] image_results[getattr(img_file, 'name', 'Unknown')] = r # 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 hasattr(r, 'boxes') and r.boxes is not None and len(r.boxes) > 0: detection_count = len(r.boxes) 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 metrics.append({ "Image": getattr(img_file, 'name', 'Unknown'), "Inference Time (ms)": inference_time if inference_time is not None else "N/A", "Detections": detection_count, "Average Confidence": f"{avg_conf:.2f}" }) eta_placeholder.empty() # 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") 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" ) # 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.0, max_value=1.0, value=0.45, # Default NMS value step=0.05, key="single_model_iou_threshold", 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 st.subheader("Annotated Images") for img_name, r in st.session_state.single_model_results.items(): try: # 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} (Conf: {conf_threshold:.2f}, IoU: {iou_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(): """ Multi-model, multi-image comparison subpage, with Weighted Scoring that uses a reciprocal-based speed metric and a real-time progress bar + ETA display. """ st.header("YOLO Models Comparison Tool (Multi-Image, Weighted Score + Progress Bar)") st.write( "Upload **one or more images** and **multiple YOLO model (.pt) files**. " "Then click **Submit** to run inference across all images with each model. " "We aggregate metrics (Avg Inference Time, Total Detections, Avg Confidence) " "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) # Example weights. You can expose them as sliders if you want user customization. alpha_detection = 0.4 beta_confidence = 0.3 gamma_speed = 0.3 # speed = reciprocal of time if st.button("Submit (Multi-Model Comparison)"): if not images or not model_files: st.error("Please upload at least one image and at least one model.") return # Initialize progress tracking total_inferences = len(images) * len(model_files) if total_inferences == 0: st.error("No valid images or models to process.") return progress_bar = st.progress(0) eta_placeholder = st.empty() start_time = time.time() steps_done = 0 # We'll store aggregated metrics here model_agg_data = {} # We'll store results for each (model, image) so we can display side-by-side model_image_results = {m.name: {} for m in model_files} for model_file in model_files: model_path = save_uploaded_file(model_file) try: model = YOLO(model_path) except Exception as e: st.error(f"Error loading model {model_file.name}: {e}") continue total_inference_time = 0.0 total_detections = 0 sum_confidences = 0.0 total_conf_count = 0 for img_file in images: # Update progress/ETA before processing next image steps_done += 1 fraction_done = steps_done / total_inferences progress_bar.progress(fraction_done) elapsed_time = time.time() - start_time time_per_step = elapsed_time / steps_done remaining_steps = total_inferences - steps_done eta_seconds = remaining_steps * time_per_step eta_placeholder.info(f"Progress: {fraction_done:.1%}. ETA: ~{eta_seconds:.1f} s") # Load image try: pil_img = Image.open(img_file).convert("RGB") np_img = np.array(pil_img) except Exception as e: st.error(f"Error reading image {img_file.name}: {e}") continue # Run inference try: # 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 r = result[0] model_image_results[model_file.name][img_file.name] = r # Accumulate inference time if isinstance(r.speed, dict) and "inference" in r.speed: total_inference_time += r.speed["inference"] # Count detections & confidence 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: 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 # After all images for this model image_count = len(images) avg_inference_time = total_inference_time / image_count if image_count > 0 else float("inf") avg_confidence = sum_confidences / total_conf_count if total_conf_count > 0 else 0.0 model_agg_data[model_file.name] = { "Model File": model_file.name, "Avg Inference Time (ms)": avg_inference_time, "Total Detections": total_detections, "Average Confidence": avg_confidence } if not model_agg_data: st.write("No valid models processed.") return # Now that all inferences are done, remove the ETA info eta_placeholder.empty() # Display aggregated metrics df = pd.DataFrame(model_agg_data.values()) # Weighted Scoring with reciprocal-based speed detection_max = df["Total Detections"].max() confidence_max = df["Average Confidence"].max() if detection_max == 0: detection_max = 1 if confidence_max == 0: confidence_max = 1 df["Detection Norm"] = df["Total Detections"] / detection_max df["Confidence Norm"] = df["Average Confidence"] / confidence_max # Convert time to speed = 1 / time, then normalize eps = 1e-9 df["Speed Val"] = 1.0 / (df["Avg Inference Time (ms)"] + eps) max_speed_val = df["Speed Val"].max() if not df["Speed Val"].isnull().all() else 1 if max_speed_val == 0: max_speed_val = 1 df["Speed Norm"] = df["Speed Val"] / max_speed_val df["Weighted Score"] = ( alpha_detection * df["Detection Norm"] + beta_confidence * df["Confidence Norm"] + 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[[ "Model File", "Avg Inference Time (ms)", "Total Detections", "Average Confidence", "Detection Norm", "Confidence Norm", "Speed Val", "Speed Norm", "Weighted Score" ]], use_container_width=True) st.markdown(f""" **Best Overall Model** based on Weighted Score: **{best_model}** (Score: {best_score:.3f}). ### Interpretation: - **Detection Norm** → fraction of the best detection count. - **Confidence Norm** → fraction of the highest average confidence. - **Speed Norm** → fraction of the highest (1/time). The fastest model is near 1; others are a fraction of that speed. If you find one factor more important, adjust the weights: - Increase **Detection** weight if you care about finding as many objects as possible. - Increase **Confidence** weight if you only trust high‐confidence detections. - 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" ) # 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.0, max_value=1.0, value=0.45, # Default NMS value step=0.05, key="multi_model_iou_threshold", 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) st.subheader("Annotated Images Grid (Row = Image, Column = Model)") 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 = 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: # 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} (Conf: {comp_conf_threshold:.2f}, IoU: {comp_iou_threshold:.2f}, Det: {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") page = st.sidebar.radio("Go to", ("YOLO Model Inference Tool", "YOLO Models Comparison Tool")) if page == "YOLO Model Inference Tool": yolo_inference_tool() else: yolo_model_comparison_tool() if __name__ == "__main__": main()