app.py

import os
import cv2
import torch
import numpy as np
import gradio as gr
from PIL import Image
import matplotlib.pyplot as plt
from transformers import AutoModel, AutoProcessor
from ultralytics import YOLO

# Custom CSS for shadcn/Radix UI inspired look
custom_css = """
:root {
  --primary: #0f172a;
  --primary-foreground: #f8fafc;
  --background: #f8fafc;
  --card: #ffffff;
  --card-foreground: #0f172a;
  --border: #e2e8f0;
  --ring: #94a3b8;
  --radius: 0.5rem;
}

.dark {
  --primary: #f8fafc;
  --primary-foreground: #0f172a;
  --background: #0f172a;
  --card: #1e293b;
  --card-foreground: #f8fafc;
  --border: #334155;
  --ring: #94a3b8;
}

.gradio-container {
  margin: 0 !important;
  padding: 0 !important;
  max-width: 100% !important;
}

.main-container {
  background-color: var(--background);
  border-radius: var(--radius);
  padding: 1.5rem;
}

.header {
  margin-bottom: 1.5rem;
  border-bottom: 1px solid var(--border);
  padding-bottom: 1rem;
}

.header h1 {
  font-size: 1.875rem;
  font-weight: 700;
  color: var(--primary);
  margin-bottom: 0.5rem;
}

.header p {
  color: var(--card-foreground);
  opacity: 0.8;
}

.tab-nav {
  background-color: var(--card);
  border: 1px solid var(--border);
  border-radius: var(--radius);
  padding: 0.25rem;
  margin-bottom: 1.5rem;
}

.tab-nav button {
  border-radius: calc(var(--radius) - 0.25rem) !important;
  font-weight: 500 !important;
  transition: all 0.2s ease-in-out !important;
}

.tab-nav button.selected {
  background-color: var(--primary) !important;
  color: var(--primary-foreground) !important;
}

.input-panel, .output-panel {
  background-color: var(--card);
  border: 1px solid var(--border);
  border-radius: var(--radius);
  padding: 1.5rem;
  box-shadow: 0 1px 3px rgba(0, 0, 0, 0.05);
}

.gr-button-primary {
  background-color: var(--primary) !important;
  color: var(--primary-foreground) !important;
  border-radius: var(--radius) !important;
  font-weight: 500 !important;
  transition: all 0.2s ease-in-out !important;
}

.gr-button-primary:hover {
  opacity: 0.9 !important;
}

.gr-form {
  border: none !important;
  background: transparent !important;
}

.gr-input, .gr-select {
  border: 1px solid var(--border) !important;
  border-radius: var(--radius) !important;
  padding: 0.5rem 0.75rem !important;
}

.gr-panel {
  border: none !important;
}

.footer {
  margin-top: 1.5rem;
  border-top: 1px solid var(--border);
  padding-top: 1rem;
  font-size: 0.875rem;
  color: var(--card-foreground);
  opacity: 0.7;
}
"""

# Available model sizes
DETECTION_MODELS = {
    "tiny": "yoloworld-t",
    "small": "yoloworld-s",
    "base": "yoloworld-b",
    "large": "yoloworld-l",
}

SEGMENTATION_MODELS = {
    "YOLOv8 Nano": "yolov8n-seg.pt",
    "YOLOv8 Small": "yolov8s-seg.pt",
    "YOLOv8 Medium": "yolov8m-seg.pt",
    "YOLOv8 Large": "yolov8l-seg.pt",
}

class YOLOWorldDetector:
    def __init__(self, model_size="base"):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model_size = model_size
        self.model_name = DETECTION_MODELS[model_size]
        
        print(f"Loading {self.model_name} on {self.device}...")
        self.model = AutoModel.from_pretrained(f"deepdatacloud/{self.model_name}", 
                                              trust_remote_code=True)
        self.model.to(self.device)
        self.processor = AutoProcessor.from_pretrained(f"deepdatacloud/{self.model_name}")
        print("Model loaded successfully!")
        
        # Segmentation models
        self.seg_models = {}

    def change_model(self, model_size):
        if model_size != self.model_size:
            self.model_size = model_size
            self.model_name = DETECTION_MODELS[model_size]
            
            print(f"Loading {self.model_name} on {self.device}...")
            self.model = AutoModel.from_pretrained(f"deepdatacloud/{self.model_name}", 
                                                  trust_remote_code=True)
            self.model.to(self.device)
            self.processor = AutoProcessor.from_pretrained(f"deepdatacloud/{self.model_name}")
            print("Model loaded successfully!")
        return f"Using {self.model_name} model"
    
    def load_seg_model(self, model_name):
        if model_name not in self.seg_models:
            print(f"Loading segmentation model {model_name}...")
            self.seg_models[model_name] = YOLO(SEGMENTATION_MODELS[model_name])
            print(f"Segmentation model {model_name} loaded successfully!")
        return self.seg_models[model_name]

    def detect(self, image, text_prompt, confidence_threshold=0.3):
        if image is None:
            return None, "No image provided"
        
        if isinstance(image, str):
            image = Image.open(image).convert("RGB")
        elif isinstance(image, np.ndarray):
            image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            
        # Process inputs
        inputs = self.processor(text=text_prompt, images=image, return_tensors="pt")
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # Run inference
        with torch.no_grad():
            outputs = self.model(**inputs)
        
        # Process results
        target_sizes = torch.tensor([image.size[::-1]], device=self.device)
        results = self.processor.post_process_object_detection(
            outputs=outputs, 
            target_sizes=target_sizes, 
            threshold=confidence_threshold
        )[0]
        
        # Convert image to numpy for drawing
        image_np = np.array(image)
        
        # Draw bounding boxes
        for box, score, label in zip(results["boxes"], results["scores"], results["labels"]):
            box = box.cpu().numpy().astype(int)
            score = score.cpu().item()
            label = label.cpu().item()
            
            # Get class name from model's config
            class_name = f"{text_prompt.split(',')[label] if label < len(text_prompt.split(',')) else 'Object'}: {score:.2f}"
            
            # Draw rectangle
            cv2.rectangle(
                image_np, 
                (box[0], box[1]), 
                (box[2], box[3]), 
                (0, 255, 0), 
                2
            )
            
            # Draw label background
            text_size = cv2.getTextSize(class_name, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]
            cv2.rectangle(
                image_np, 
                (box[0], box[1] - text_size[1] - 5), 
                (box[0] + text_size[0], box[1]), 
                (0, 255, 0), 
                -1
            )
            
            # Draw text
            cv2.putText(
                image_np, 
                class_name, 
                (box[0], box[1] - 5), 
                cv2.FONT_HERSHEY_SIMPLEX, 
                0.5, 
                (0, 0, 0), 
                2
            )
        
        # Convert results to JSON format (percentages)
        json_results = []
        img_height, img_width = image_np.shape[:2]
        
        for box, score, label in zip(results["boxes"], results["scores"], results["labels"]):
            box = box.cpu().numpy()
            x1, y1, x2, y2 = box
            
            json_results.append({
                "bbox": {
                    "x": (x1 / img_width) * 100,
                    "y": (y1 / img_height) * 100,
                    "width": ((x2 - x1) / img_width) * 100,
                    "height": ((y2 - y1) / img_height) * 100
                },
                "score": float(score.cpu().item()),
                "label": int(label.cpu().item()),
                "label_text": text_prompt.split(',')[label] if label < len(text_prompt.split(',')) else 'Object'
            })
        
        return cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR), json_results
    
    def segment(self, image, model_name, confidence_threshold=0.3):
        if image is None:
            return None, "No image provided"
        
        # Load segmentation model if not already loaded
        model = self.load_seg_model(model_name)
        
        # Run inference
        results = model(image, conf=confidence_threshold)
        
        # Create visualization
        fig, ax = plt.subplots(1, 1, figsize=(12, 9))
        ax.axis('off')
        
        # Plot segmentation results
        res_plotted = results[0].plot()
        
        # Convert results to JSON format (percentages)
        json_results = []
        if hasattr(results[0], 'masks') and results[0].masks is not None:
            img_height, img_width = results[0].orig_shape
            
            for i, (box, mask, cls, conf) in enumerate(zip(
                results[0].boxes.xyxy.cpu().numpy(),
                results[0].masks.data.cpu().numpy(),
                results[0].boxes.cls.cpu().numpy(),
                results[0].boxes.conf.cpu().numpy()
            )):
                x1, y1, x2, y2 = box
                
                # Convert mask to polygon for SVG-like representation
                # Simplified approach - in production you might want a more sophisticated polygon extraction
                contours, _ = cv2.findContours((mask > 0.5).astype(np.uint8), 
                                              cv2.RETR_EXTERNAL, 
                                              cv2.CHAIN_APPROX_SIMPLE)
                
                if contours:
                    # Get the largest contour
                    largest_contour = max(contours, key=cv2.contourArea)
                    # Simplify the contour
                    epsilon = 0.005 * cv2.arcLength(largest_contour, True)
                    approx = cv2.approxPolyDP(largest_contour, epsilon, True)
                    
                    # Convert to percentage coordinates
                    points = []
                    for point in approx:
                        x, y = point[0]
                        points.append({
                            "x": (x / img_width) * 100,
                            "y": (y / img_height) * 100
                        })
                    
                    json_results.append({
                        "bbox": {
                            "x": (x1 / img_width) * 100,
                            "y": (y1 / img_height) * 100,
                            "width": ((x2 - x1) / img_width) * 100,
                            "height": ((y2 - y1) / img_height) * 100
                        },
                        "score": float(conf),
                        "label": int(cls),
                        "label_text": results[0].names[int(cls)],
                        "polygon": points
                    })
        
        return res_plotted, json_results

# Initialize detector with default model
detector = YOLOWorldDetector(model_size="base")

def detection_inference(image, text_prompt, confidence, model_size):
    # Update model if needed
    detector.change_model(model_size)
    
    # Run detection
    result_image, json_results = detector.detect(
        image, 
        text_prompt, 
        confidence_threshold=confidence
    )
    
    return result_image, str(json_results)

def segmentation_inference(image, confidence, model_name):
    # Run segmentation
    result_image, json_results = detector.segment(
        image,
        model_name,
        confidence_threshold=confidence
    )
    
    return result_image, str(json_results)

# Create Gradio interface
with gr.Blocks(title="YOLO Vision Suite", css=custom_css) as demo:
    with gr.Column(elem_classes="main-container"):
        with gr.Column(elem_classes="header"):
            gr.Markdown("# YOLO Vision Suite")
            gr.Markdown("Advanced object detection and segmentation powered by YOLO models")
        
        with gr.Tabs(elem_classes="tab-nav") as tabs:
            with gr.TabItem("Object Detection", elem_id="detection-tab"):
                with gr.Row():
                    with gr.Column(elem_classes="input-panel"):
                        gr.Markdown("### Input")
                        input_image = gr.Image(label="Upload Image", type="numpy")
                        text_prompt = gr.Textbox(
                            label="Text Prompt", 
                            placeholder="person, car, dog",
                            value="person, car, dog",
                            elem_classes="gr-input"
                        )
                        with gr.Row():
                            confidence = gr.Slider(
                                minimum=0.1, 
                                maximum=1.0, 
                                value=0.3, 
                                step=0.05, 
                                label="Confidence Threshold"
                            )
                            model_dropdown = gr.Dropdown(
                                choices=list(DETECTION_MODELS.keys()),
                                value="base",
                                label="Model Size",
                                elem_classes="gr-select"
                            )
                        detect_button = gr.Button("Detect Objects", elem_classes="gr-button-primary")
                    
                    with gr.Column(elem_classes="output-panel"):
                        gr.Markdown("### Results")
                        output_image = gr.Image(label="Detection Result")
                        with gr.Accordion("JSON Output", open=False):
                            json_output = gr.Textbox(
                                label="Bounding Box Data (Percentage Coordinates)",
                                elem_classes="gr-input"
                            )
            
            with gr.TabItem("Segmentation", elem_id="segmentation-tab"):
                with gr.Row():
                    with gr.Column(elem_classes="input-panel"):
                        gr.Markdown("### Input")
                        seg_input_image = gr.Image(label="Upload Image", type="numpy")
                        with gr.Row():
                            seg_confidence = gr.Slider(
                                minimum=0.1, 
                                maximum=1.0, 
                                value=0.3, 
                                step=0.05, 
                                label="Confidence Threshold"
                            )
                            seg_model_dropdown = gr.Dropdown(
                                choices=list(SEGMENTATION_MODELS.keys()),
                                value="YOLOv8 Small",
                                label="Model Size",
                                elem_classes="gr-select"
                            )
                        segment_button = gr.Button("Segment Image", elem_classes="gr-button-primary")
                    
                    with gr.Column(elem_classes="output-panel"):
                        gr.Markdown("### Results")
                        seg_output_image = gr.Image(label="Segmentation Result")
                        with gr.Accordion("JSON Output", open=False):
                            seg_json_output = gr.Textbox(
                                label="Segmentation Data (Percentage Coordinates)",
                                elem_classes="gr-input"
                            )
        
        with gr.Column(elem_classes="footer"):
            gr.Markdown("""
            ### Tips
            - For object detection, enter comma-separated text prompts to specify what to detect
            - For segmentation, the model will identify common objects automatically
            - Larger models provide better accuracy but require more processing power
            - The JSON output provides coordinates as percentages of image dimensions, compatible with SVG
            """)
    
    # Set up event handlers
    detect_button.click(
        detection_inference,
        inputs=[input_image, text_prompt, confidence, model_dropdown],
        outputs=[output_image, json_output]
    )
    
    segment_button.click(
        segmentation_inference,
        inputs=[seg_input_image, seg_confidence, seg_model_dropdown],
        outputs=[seg_output_image, seg_json_output]
    )

if __name__ == "__main__":
    demo.launch()