Update app.py

app.py

@@ -1,164 +1,222 @@
-from flask import Flask, request, jsonify
-from flask_cors import CORS
-import cv2
-import numpy as np
-import base64
-import io
-from collections import Counter
-from PIL import Image
-from ultralytics import YOLO
-import os
-
-# App Configuration
-app = Flask(__name__)
-CORS(app)
-
-# Model paths
-MODELS = {
-    'yolov8s.pt': './model/yolov8s.pt',
-    'yolov9m.pt': './model/yolov9m.pt'
-}
-
-# Load models on demand
-loaded_models = {}
-
-def get_model(model_name):
-    """Load model if not already loaded"""
-    if model_name not in loaded_models:
-        if model_name in MODELS and os.path.exists(MODELS[model_name]):
-            loaded_models[model_name] = YOLO(MODELS[model_name])
-        else:
-            raise ValueError(f"Model {model_name} not found")
-    return loaded_models[model_name]
-
-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)
-
-@app.route('/detect', methods=['POST'])
-def detect_objects():
-    try:
-        # Image receive karna
-        image_base64 = request.json.get('image', '')
-        
-        # Get selected model
-        model_name = request.json.get('model', 'yolov9m.pt')
-        model = get_model(model_name)
-        
-        # Get confidence threshold
-        confidence = request.json.get('confidence', 0.25)
-        
-        # Image decode karna
-        image = decode_base64_image(image_base64)
-        
-        # Object detection
-        results = model(image, conf=confidence)
-        
-        # Detected objects ko process karna
-        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)
-        
-        # Object grouping
-        object_counts = Counter(det['class'] for det in detections)
-        grouped_objects = [
-            {'class': obj, 'count': count} 
-            for obj, count in object_counts.items()
-        ]
-        
-        return jsonify({
-            'detections': detections,
-            'grouped_objects': grouped_objects,
-            'model_used': model_name
-        })
-    
-    except Exception as e:
-        return jsonify({'error': str(e)}), 500
-
-@app.route('/available-models', methods=['GET'])
-def get_available_models():
-    """Available object detection models"""
-    return jsonify({
-        'models': [
-            {'name': 'yolov8s.pt', 'type': 'Object Detection', 'description': 'YOLOv8s (Fastest)'},
-            {'name': 'yolov9m.pt', 'type': 'Object Detection', 'description': 'YOLOv9m (Highest Accuracy)'},
-        ]
-    })
-
-@app.route('/detect-multiple', methods=['POST'])
-def detect_multiple_objects():
-    """Multiple images ke liye detection support"""
-    try:
-        images_base64 = request.json.get('images', [])
-        
-        # Get selected model
-        model_name = request.json.get('model', 'yolov9m.pt')
-        model = get_model(model_name)
-        
-        # Get confidence threshold
-        confidence = request.json.get('confidence', 0.25)
-        
-        all_detections = []
-        
-        for image_base64 in images_base64:
-            # Image decode karna
-            image = decode_base64_image(image_base64)
-            
-            # Object detection
-            results = model(image, conf=confidence)
-            
-            # Detected objects ko process karna
-            image_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)
-                    }
-                    image_detections.append(detection)
-            
-            all_detections.append(image_detections)
-        
-        return jsonify({
-            'detections': all_detections,
-            'model_used': model_name
-        })
-    
-    except Exception as e:
-        return jsonify({'error': str(e)}), 500
-
-if __name__ == '__main__':
-    app.run(debug=True, port=5000)
+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
+""")