import streamlit as st import cv2 import numpy as np import base64 import io from collections import Counter from PIL import Image from ultralytics import YOLO import os # Set page config st.set_page_config(page_title="Object Detection App", layout="wide") # Model paths MODELS = { 'yolov8s.pt': './model/yolov8s.pt', 'yolov9m.pt': './model/yolov9m.pt' } # Load models on demand @st.cache_resource def get_model(model_name): """Load model if not already loaded""" if model_name in MODELS and os.path.exists(MODELS[model_name]): return YOLO(MODELS[model_name]) else: raise ValueError(f"Model {model_name} not found") def decode_base64_image(base64_string): """Base64 image string ko decode karna""" # Remove data URL prefix if present if ',' in base64_string: base64_string = base64_string.split(',')[1] image_data = base64.b64decode(base64_string) image = Image.open(io.BytesIO(image_data)) return np.array(image) def process_detections(results, model): """Process detection results into standard format""" detections = [] for result in results: boxes = result.boxes for box in boxes: # Bounding box coordinates x1, y1, x2, y2 = box.xyxy[0] # Confidence aur class conf = box.conf[0] cls = int(box.cls[0]) class_name = model.names[cls] # Detection object banana detection = { 'bbox': [float(x1), float(y1), float(x2-x1), float(y2-y1)], 'class': class_name, 'confidence': float(conf) } detections.append(detection) return detections # App title st.title("Object Detection App") # Sidebar for settings st.sidebar.title("Settings") # Available models info available_models = [ {'name': 'yolov8s.pt', 'type': 'Object Detection', 'description': 'YOLOv8s (Fastest)'}, {'name': 'yolov9m.pt', 'type': 'Object Detection', 'description': 'YOLOv9m (Highest Accuracy)'}, ] # Model selection model_options = {m['name']: f"{m['name']} - {m['description']}" for m in available_models} model_name = st.sidebar.selectbox("Select Model", options=list(model_options.keys()), format_func=lambda x: model_options[x]) # Confidence threshold confidence = st.sidebar.slider("Confidence Threshold", min_value=0.1, max_value=1.0, value=0.25, step=0.05) # Tab selection tab1, tab2 = st.tabs(["Single Image", "Multiple Images"]) with tab1: st.header("Single Image Detection") # Image upload uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"]) if uploaded_file is not None: # Display uploaded image image = Image.open(uploaded_file) st.image(image, caption="Uploaded Image", use_column_width=True) # Process button if st.button("Detect Objects"): try: with st.spinner("Detecting objects..."): # Load model model = get_model(model_name) # Convert to numpy array image_np = np.array(image) # Object detection results = model(image_np, conf=confidence) # Process results detections = process_detections(results, model) # Object grouping object_counts = Counter(det['class'] for det in detections) grouped_objects = [ {'class': obj, 'count': count} for obj, count in object_counts.items() ] # Display results if any detections found if detections: # Draw bounding boxes on image result_image = image_np.copy() for det in detections: x, y, w, h = [int(val) for val in det['bbox']] cv2.rectangle(result_image, (x, y), (x+w, y+h), (0, 255, 0), 2) cv2.putText(result_image, f"{det['class']} {det['confidence']:.2f}", (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) # Show image with detections st.image(result_image, caption="Detection Results", use_column_width=True) # Display summary st.subheader("Detection Summary") for obj in grouped_objects: st.write(f"- {obj['class']}: {obj['count']}") # Display detection details st.subheader("Detection Details") for i, det in enumerate(detections, 1): st.write(f"#{i}: {det['class']} (Confidence: {det['confidence']:.2f})") else: st.info("No objects detected in the image.") except Exception as e: st.error(f"Error processing image: {str(e)}") with tab2: st.header("Multiple Images Detection") uploaded_files = st.file_uploader("Upload multiple images", type=["jpg", "jpeg", "png"], accept_multiple_files=True) if uploaded_files: st.write(f"{len(uploaded_files)} images uploaded") # Process button if st.button("Detect Objects in All Images"): try: with st.spinner("Detecting objects in multiple images..."): # Load model model = get_model(model_name) # Process each image all_detections = [] for i, file in enumerate(uploaded_files): # Read image image = Image.open(file) image_np = np.array(image) # Object detection results = model(image_np, conf=confidence) # Process results detections = process_detections(results, model) all_detections.append(detections) # Create columns for image display col1, col2 = st.columns(2) with col1: st.write(f"Image {i+1}: {file.name}") st.image(image, caption=f"Original - {file.name}", use_column_width=True) with col2: # Draw bounding boxes result_image = image_np.copy() for det in detections: x, y, w, h = [int(val) for val in det['bbox']] cv2.rectangle(result_image, (x, y), (x+w, y+h), (0, 255, 0), 2) cv2.putText(result_image, f"{det['class']} {det['confidence']:.2f}", (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) st.image(result_image, caption=f"Detections - {file.name}", use_column_width=True) # Display detections for this image object_counts = Counter(det['class'] for det in detections) st.write("Detected objects:") for obj, count in object_counts.items(): st.write(f"- {obj}: {count}") st.divider() # Overall summary st.subheader("Overall Detection Summary") all_objects = [] for detections in all_detections: all_objects.extend([det['class'] for det in detections]) total_counts = Counter(all_objects) for obj, count in total_counts.items(): st.write(f"- {obj}: {count} (across all images)") except Exception as e: st.error(f"Error processing images: {str(e)}") # About section st.sidebar.markdown("---") st.sidebar.header("About") st.sidebar.info(""" This app uses YOLO models for object detection. - YOLOv8s: Faster detection - YOLOv9m: Higher accuracy """)