Update app.py

app.py

@@ -1,305 +1,122 @@
-import os
 import gradio as gr
-import requests
-from PIL import Image, ImageDraw, ImageFont
-import numpy as np
 import cv2
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+from ultralytics import YOLO
 import time
 
-# Groq API credentials and endpoint
-API_KEY = "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"
-API_URL = "https://api.groq.com/openai/v1/chat/completions"
-
-# Performance metrics
-startup_time = time.time()
-processed_images = 0
-
-# Model configuration
-MODEL_NAMES = {
-    "nano": "yolov8n.pt",
-    "small": "yolov8s.pt",
-    "medium": "yolov8m.pt",
-    "large": "yolov8l.pt",
-    "extra large": "yolov8x.pt"
-}
-
-class YOLOModel:
-    def __init__(self):
-        self.model = None
-        self.model_size = None
-        self.last_used = None
-        
-    def load_model(self, model_size="medium"):
-        global model_loaded
-        try:
-            from ultralytics import YOLO
-            model_path = MODEL_NAMES.get(model_size, "yolov8m.pt")
-            self.model = YOLO(model_path)
-            self.model_size = model_size
-            self.last_used = time.time()
-            model_loaded = True
-            return f"YOLOv8 {model_size} model loaded successfully", True
-        except Exception as e:
-            return f"Error loading YOLO model: {str(e)}", False
-    
-    def get_model(self):
-        if self.model is None:
-            return None
-        self.last_used = time.time()
-        return self.model
-
-yolo_model = YOLOModel()
-
-def get_system_health():
-    """Get system health metrics"""
-    uptime = time.time() - startup_time
-    return {
-        "uptime": f"{uptime:.1f} seconds",
-        "images_processed": processed_images,
-        "model_status": "Loaded" if yolo_model.model else "Not loaded",
-        "model_size": yolo_model.model_size if yolo_model.model else "None",
-        "last_used": f"{time.time() - yolo_model.last_used:.1f} sec ago" if yolo_model.last_used else "Never"
+# Initialize YOLO model
+def load_model(model_size="medium"):
+    model_map = {
+        "nano": "yolov8n.pt",
+        "small": "yolov8s.pt",
+        "medium": "yolov8m.pt",
+        "large": "yolov8l.pt",
+        "xlarge": "yolov8x.pt"
     }
-
-def groq_chat_api(user_message, history=None):
-    """Function to call Groq API for chat completions"""
-    if not user_message or user_message.strip() == "":
-        return "Please enter a question to get a response."
-        
     try:
-        messages = [{"role": "user", "content": user_message}]
-        if history:
-            # Include conversation history if provided
-            for user_msg, bot_response in history:
-                messages.insert(0, {"role": "assistant", "content": bot_response})
-                messages.insert(0, {"role": "user", "content": user_msg})
-                
-        payload = {
-            "messages": messages,
-            "model": "llama3-70b-8192",
-            "temperature": 0.7,
-            "max_tokens": 2048,
-            "top_p": 0.9,
-        }
-
-        headers = {
-            "Authorization": f"Bearer {API_KEY}",
-            "Content-Type": "application/json"
-        }
-
-        response = requests.post(API_URL, json=payload, headers=headers, timeout=30)
-
-        if response.status_code == 200:
-            return response.json()["choices"][0]["message"]["content"]
-        else:
-            return f"Error: Failed to get response (HTTP {response.status_code}). Please try again."
-
-    except requests.Timeout:
-        return "Error: Request timed out. The server took too long to respond."
+        model = YOLO(model_map[model_size])
+        return model, f"✅ {model_size.capitalize()} model loaded successfully"
     except Exception as e:
-        return f"Error processing request: {str(e)}"
+        return None, f"❌ Error loading model: {str(e)}"
 
-def draw_detection_info(image, detections):
-    """Draw detection information directly on the image"""
-    try:
-        # Convert PIL Image to OpenCV format if needed
-        if isinstance(image, Image.Image):
-            image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-        
-        # Create a blank space at the bottom for the text
-        h, w = image.shape[:2]
-        text_area = np.zeros((150, w, 3), dtype=np.uint8)
-        text_area.fill(245)  # Light gray background
-        combined = np.vstack([image, text_area])
-        
-        # Convert back to PIL Image to use PIL's text drawing
-        combined_pil = Image.fromarray(cv2.cvtColor(combined, cv2.COLOR_BGR2RGB))
-        draw = ImageDraw.Draw(combined_pil)
-        
-        try:
-            # Try to use a nicer font if available
-            font = ImageFont.truetype("arial.ttf", 14)
-        except:
-            font = ImageFont.load_default()
-        
-        # Prepare detection text
-        text_lines = ["Detection Results:"]
-        for i, det in enumerate(detections[:5]):  # Show max 5 detections
-            text_lines.append(f"{i+1}. {det['class']} (Confidence: {det['confidence']:.1%})")
-        
-        if not detections:
-            text_lines.append("No objects detected")
-        
-        # Draw text
-        y_offset = h + 10
-        for line in text_lines:
-            draw.text((10, y_offset), line, font=font, fill=(0, 0, 0))
-            y_offset += 20
-        
-        return combined_pil
+# Enhanced detection function
+def detect_objects(image, model_size="medium", conf_threshold=0.35, iou_threshold=0.45):
+    # Load model (load fresh each time to change model size dynamically)
+    model, status = load_model(model_size)
+    if model is None:
+        return image, status
     
-    except Exception as e:
-        print(f"Error in draw_detection_info: {str(e)}")
-        return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
-
-def process_image(image, conf_threshold=0.4):
-    """Process image with YOLOv8 model and return annotated image with detection results"""
-    global processed_images
+    # Convert image to numpy array if needed
+    if isinstance(image, np.ndarray):
+        img = image
+    else:
+        img = np.array(image)
     
-    if image is None:
-        return None, "Error: No image provided. Please upload an image."
+    # Run inference with improved parameters
+    results = model.predict(
+        source=img,
+        conf=conf_threshold,
+        iou=iou_threshold,
+        imgsz=640,  # Increased from default 320 for better detection
+        augment=True,  # Enable test-time augmentation
+        visualize=False  # Disable visualization for speed
+    )
     
-    model = yolo_model.get_model()
-    if model is None:
-        return image, "Error: YOLO model not loaded. Please initialize the model first."
+    # Extract detailed detection information
+    detections = []
+    for r in results:
+        boxes = r.boxes.cpu().numpy()
+        for i, box in enumerate(boxes):
+            detections.append({
+                "class": model.names[int(box.cls[0])],
+                "confidence": float(box.conf[0]),
+                "bbox": box.xyxy[0].tolist(),
+                "id": i+1
+            })
     
-    try:
-        # Convert to numpy array if it's not already
-        if isinstance(image, Image.Image):
-            img_array = np.array(image)
-        else:
-            img_array = image
-        
-        # Run YOLOv8 inference
-        results = model.predict(img_array, conf=conf_threshold)
-        processed_images += 1
-        
-        # Get the annotated image
-        annotated_img = results[0].plot()
-        
-        # Get detailed detection results
-        detections = []
-        for r in results:
-            boxes = r.boxes  # Detected objects
-            for box in boxes:
-                class_id = int(box.cls[0])
-                class_name = model.names[class_id]
-                confidence = float(box.conf[0])
-                detections.append({
-                    "class": class_name,
-                    "confidence": confidence,
-                    "bbox": box.xyxy[0].tolist()
-                })
-        
-        # Sort detections by confidence (highest first)
-        detections.sort(key=lambda x: x["confidence"], reverse=True)
-        
-        # Create detection text
-        if not detections:
-            detection_text = "No objects detected with current confidence threshold"
-        else:
-            detection_text = "Detected objects:\n"
-            for i, det in enumerate(detections[:10]):  # Show max 10 detections
-                detection_text += f"{i+1}. {det['class']} (Confidence: {det['confidence']:.1%})\n"
-        
-        # Draw information on the image itself
-        final_image = draw_detection_info(annotated_img, detections)
-        
-        return final_image, detection_text
-    except Exception as e:
-        return image, f"Error processing image: {str(e)}"
-
-# Theme colors
-theme = gr.themes.Default(
-    primary_hue="blue",
-    secondary_hue="gray",
-    font=[gr.themes.GoogleFont('Montserrat'), 'Arial', 'sans-serif']
-)
+    # Sort detections by confidence (highest first)
+    detections.sort(key=lambda x: x["confidence"], reverse=True)
+    
+    # Create annotated image with enhanced visualization
+    annotated_img = results[0].plot(
+        line_width=2,
+        font_size=14,
+        pil=True
+    )
+    
+    # Generate detailed detection text
+    detection_text = "Detection Results:\n"
+    if not detections:
+        detection_text = "No objects detected"
+    else:
+        for det in detections[:20]:  # Show top 20 detections
+            detection_text += (
+                f"{det['id']}. {det['class'].upper()} "
+                f"(Confidence: {det['confidence']:.1%})\n"
+                f"   Bounding Box: {[int(x) for x in det['bbox']]}\n"
+            )
+    
+    return annotated_img, detection_text
 
-# Create Gradio interface
-with gr.Blocks(title="Enhanced AI Vision & Chat Assistant", theme=theme) as demo:
-    gr.Markdown("""
-    # 🚀 Enhanced AI Vision & Chat Assistant
-    ### Powered by YOLOv8 Object Detection and Groq's Llama 3 LLM
-    """)
+# Gradio interface
+with gr.Blocks(title="Enhanced Object Detector") as demo:
+    gr.Markdown("# 🚀 Enhanced Object Detection with YOLOv8")
+    gr.Markdown("### Get comprehensive detections with adjustable parameters")
     
-    with gr.Tab("🖼️ Object Detection"):
-        with gr.Row():
-            with gr.Column():
-                gr.Markdown("### Image Settings")
-                image_input = gr.Image(label="Upload Image", type="numpy")
-                conf_slider = gr.Slider(0.1, 0.9, value=0.4, label="Confidence Threshold")
-                load_model_dropdown = gr.Dropdown(
-                    list(MODEL_NAMES.keys()),
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(label="Input Image", type="pil")
+            with gr.Accordion("Detection Settings", open=False):
+                model_size = gr.Dropdown(
+                    ["nano", "small", "medium", "large", "xlarge"],
                     value="medium",
-                    label="Select YOLO Model Size"
+                    label="Model Size"
                 )
-                with gr.Row():
-                    load_model_button = gr.Button("🚀 Initialize Model")
-                    detect_button = gr.Button("🔍 Detect Objects")
-                
-            with gr.Column():
-                gr.Markdown("### Detection Results")
-                image_output = gr.Image(label="Annotated Image")
-                text_output = gr.Textbox(label="Detailed Detection", lines=8)
-        
-        load_model_button.click(
-            fn=lambda size: yolo_model.load_model(size),
-            inputs=[load_model_dropdown],
-            outputs=[text_output]
-        )
-        
-        detect_button.click(
-            fn=process_image,
-            inputs=[image_input, conf_slider],
-            outputs=[image_output, text_output]
-        )
-    
-    with gr.Tab("💬 AI Chat Assistant"):
-        gr.Markdown("### Chat with Llama 3 (70B) via Groq")
-        chatbot = gr.Chatbot(height=400, label="Conversation History")
-        with gr.Row():
-            chat_input = gr.Textbox(label="Your Message", placeholder="Ask me anything...", lines=2)
-            chat_button = gr.Button("✉️ Send")
-                
-        clear_button = gr.Button("🧹 Clear Chat")
-        
-        chat_button.click(
-            fn=groq_chat_api,
-            inputs=[chat_input, chatbot],
-            outputs=[chat_input]
-        ).then(
-            lambda user_msg, chatbot: chatbot + [[user_msg, None]],
-            inputs=[chat_input, chatbot],
-            outputs=[chatbot],
-            queue=False
-        ).then(
-            lambda bot_msg, chatbot: (
-                None,  # Clear input
-                chatbot[:-1] + [(chatbot[-1][0], bot_msg)]  # Update last message
-            ),
-            inputs=[chat_input, chatbot],
-            outputs=[chat_input, chatbot]
-        )
-        
-        clear_button.click(
-            lambda: None,
-            inputs=None,
-            outputs=[chatbot],
-            queue=False
-        )
+                conf_slider = gr.Slider(
+                    minimum=0.1, maximum=0.9, step=0.05,
+                    value=0.35, label="Confidence Threshold"
+                )
+                iou_slider = gr.Slider(
+                    minimum=0.1, maximum=0.9, step=0.05,
+                    value=0.45, label="IOU Threshold"
+                )
+            detect_btn = gr.Button("🔍 Detect Objects", variant="primary")
+
+        with gr.Column():
+            image_output = gr.Image(label="Detected Objects", interactive=False)
+            text_output = gr.Textbox(
+                label="Detection Details",
+                lines=15,
+                max_lines=20,
+                interactive=False
+            )
     
-    with gr.Tab("⚙️ System Health"):
-        gr.Markdown("### Performance Metrics")
-        health_button = gr.Button("🔄 Refresh Health Status")
-        health_output = gr.JSON(label="System Metrics")
-        
-        health_button.click(
-            fn=get_system_health,
-            inputs=None,
-            outputs=health_output
-        )
-        
-        gr.Examples(
-            examples=[
-                ["What is YOLOv8 and how does it work?"],
-                ["How can AI be used for industrial equipment monitoring?"],
-                ["Explain the latest advancements in computer vision"]
-            ],
-            inputs=[chat_input],
-            label="Example Questions"
-        )
+    detect_btn.click(
+        fn=detect_objects,
+        inputs=[image_input, model_size, conf_slider, iou_slider],
+        outputs=[image_output, text_output]
+    )
 
-# Launch the app
 if __name__ == "__main__":
-    demo.launch(share=False)
+    demo.launch()