app.py

import gradio as gr
import torch
from transformers import DetrImageProcessor, DetrForObjectDetection
from transformers import YolosImageProcessor, YolosForObjectDetection
from transformers import DetrForSegmentation
from PIL import Image, ImageDraw, ImageStat
import requests
from io import BytesIO
import base64
from collections import Counter
import logging
from fastapi import FastAPI, File, UploadFile, HTTPException, Form
from fastapi.responses import JSONResponse
import uvicorn
import pandas as pd
import traceback
import os

# Set up logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

# Constants
CONFIDENCE_THRESHOLD = 0.5
VALID_MODELS = [
    "facebook/detr-resnet-50",
    "facebook/detr-resnet-101",
    "facebook/detr-resnet-50-panoptic",
    "facebook/detr-resnet-101-panoptic",
    "hustvl/yolos-tiny",
    "hustvl/yolos-base"
]
MODEL_DESCRIPTIONS = {
    "facebook/detr-resnet-50": "DETR with ResNet-50 backbone for object detection. Fast and accurate for general use.",
    "facebook/detr-resnet-101": "DETR with ResNet-101 backbone for object detection. More accurate but slower than ResNet-50.",
    "facebook/detr-resnet-50-panoptic": "DETR with ResNet-50 for panoptic segmentation. Detects objects and segments scenes.",
    "facebook/detr-resnet-101-panoptic": "DETR with ResNet-101 for panoptic segmentation. High accuracy for complex scenes.",
    "hustvl/yolos-tiny": "YOLOS Tiny model. Lightweight and fast, ideal for resource-constrained environments.",
    "hustvl/yolos-base": "YOLOS Base model. Balances speed and accuracy for object detection."
}

# Lazy model loading
models = {}
processors = {}

def process(image, model_name):
    """Process an image and return detected image, objects, confidences, unique objects, unique confidences, and properties."""
    try:
        if model_name not in VALID_MODELS:
            raise ValueError(f"Invalid model: {model_name}. Choose from: {VALID_MODELS}")

        # Load model and processor
        if model_name not in models:
            logger.info(f"Loading model: {model_name}")
            if "yolos" in model_name:
                models[model_name] = YolosForObjectDetection.from_pretrained(model_name)
                processors[model_name] = YolosImageProcessor.from_pretrained(model_name)
            elif "panoptic" in model_name:
                models[model_name] = DetrForSegmentation.from_pretrained(model_name)
                processors[model_name] = DetrImageProcessor.from_pretrained(model_name)
            else:
                models[model_name] = DetrForObjectDetection.from_pretrained(model_name)
                processors[model_name] = DetrImageProcessor.from_pretrained(model_name)

        model, processor = models[model_name], processors[model_name]
        inputs = processor(images=image, return_tensors="pt")

        with torch.no_grad():
            outputs = model(**inputs)

        target_sizes = torch.tensor([image.size[::-1]])
        draw = ImageDraw.Draw(image)
        object_names = []
        confidence_scores = []
        object_counter = Counter()

        if "panoptic" in model_name:
            processed_sizes = torch.tensor([[inputs["pixel_values"].shape[2], inputs["pixel_values"].shape[3]]])
            results = processor.post_process_panoptic(outputs, target_sizes=target_sizes, processed_sizes=processed_sizes)[0]

            for segment in results["segments_info"]:
                label = segment["label_id"]
                label_name = model.config.id2label.get(label, "Unknown")
                score = segment.get("score", 1.0)

                if "masks" in results and segment["id"] < len(results["masks"]):
                    mask = results["masks"][segment["id"]].cpu().numpy()
                    if mask.shape[0] > 0 and mask.shape[1] > 0:
                        mask_image = Image.fromarray((mask * 255).astype("uint8"))
                        colored_mask = Image.new("RGBA", image.size, (0, 0, 0, 0))
                        mask_draw = ImageDraw.Draw(colored_mask)
                        r, g, b = (segment["id"] * 50) % 255, (segment["id"] * 100) % 255, (segment["id"] * 150) % 255
                        mask_draw.bitmap((0, 0), mask_image, fill=(r, g, b, 128))
                        image = Image.alpha_composite(image.convert("RGBA"), colored_mask).convert("RGB")
                        draw = ImageDraw.Draw(image)

                if score > CONFIDENCE_THRESHOLD:
                    object_names.append(label_name)
                    confidence_scores.append(float(score))
                    object_counter[label_name] = float(score)
        else:
            results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)[0]

            for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
                if score > CONFIDENCE_THRESHOLD:
                    x, y, x2, y2 = box.tolist()
                    draw.rectangle([x, y, x2, y2], outline="#32CD32", width=2)
                    label_name = model.config.id2label.get(label.item(), "Unknown")
                    # Place text at top-right corner, outside the box, with smaller size
                    text = f"{label_name}: {score:.2f}"
                    text_bbox = draw.textbbox((0, 0), text)
                    text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
                    draw.text((x2 - text_width - 2, y - text_height - 2), text, fill="#32CD32")
                    object_names.append(label_name)
                    confidence_scores.append(float(score))
                    object_counter[label_name] = float(score)

        unique_objects = list(object_counter.keys())
        unique_confidences = [object_counter[obj] for obj in unique_objects]

        # Image properties
        file_size = "Unknown"
        if hasattr(image, "fp") and image.fp is not None:
            buffered = BytesIO()
            image.save(buffered, format="PNG")
            file_size = f"{len(buffered.getvalue()) / 1024:.2f} KB"
        
        # Color statistics
        try:
            stat = ImageStat.Stat(image)
            color_stats = {
                "mean": [f"{m:.2f}" for m in stat.mean],
                "stddev": [f"{s:.2f}" for s in stat.stddev]
            }
        except Exception as e:
            logger.error(f"Error calculating color statistics: {str(e)}")
            color_stats = {"mean": "Error", "stddev": "Error"}

        properties = {
            "Format": image.format if hasattr(image, "format") and image.format else "Unknown",
            "Size": f"{image.width}x{image.height}",
            "Width": f"{image.width} px",
            "Height": f"{image.height} px",
            "Mode": image.mode,
            "Aspect Ratio": f"{round(image.width / image.height, 2) if image.height != 0 else 'Undefined'}",
            "File Size": file_size,
            "Mean (R,G,B)": ", ".join(color_stats["mean"]) if isinstance(color_stats["mean"], list) else color_stats["mean"],
            "StdDev (R,G,B)": ", ".join(color_stats["stddev"]) if isinstance(color_stats["stddev"], list) else color_stats["stddev"]
        }

        return image, object_names, confidence_scores, unique_objects, unique_confidences, properties
    except Exception as e:
        logger.error(f"Error in process: {str(e)}\n{traceback.format_exc()}")
        raise

# FastAPI Setup
app = FastAPI(title="Object Detection API")

@app.post("/detect")
async def detect_objects_endpoint(
    file: UploadFile = File(None),
    image_url: str = Form(None),
    model_name: str = Form(VALID_MODELS[0])
):
    """FastAPI endpoint to detect objects in an image from file or URL."""
    try:
        if (file is None and not image_url) or (file is not None and image_url):
            raise HTTPException(status_code=400, detail="Provide either an image file or an image URL, but not both.")

        if file:
            if not file.content_type.startswith("image/"):
                raise HTTPException(status_code=400, detail="File must be an image")
            contents = await file.read()
            image = Image.open(BytesIO(contents)).convert("RGB")
        else:
            response = requests.get(image_url, timeout=10)
            response.raise_for_status()
            image = Image.open(BytesIO(response.content)).convert("RGB")

        if model_name not in VALID_MODELS:
            raise HTTPException(status_code=400, detail=f"Invalid model. Choose from: {VALID_MODELS}")

        detected_image, detected_objects, detected_confidences, unique_objects, unique_confidences, _ = process(image, model_name)

        buffered = BytesIO()
        detected_image.save(buffered, format="PNG")
        img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
        img_url = f"data:image/png;base64,{img_base64}"

        return JSONResponse(content={
            "image_url": img_url,
            "detected_objects": detected_objects,
            "confidence_scores": detected_confidences,
            "unique_objects": unique_objects,
            "unique_confidence_scores": unique_confidences
        })
    except Exception as e:
        logger.error(f"Error in FastAPI endpoint: {str(e)}\n{traceback.format_exc()}")
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")

# Gradio UI
def create_gradio_ui():
    with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="gray")) as demo:
        gr.Markdown(
            """
            # 🚀 Object Detection App
            Upload an image or provide a URL to detect objects using state-of-the-art transformer models (DETR, YOLOS).
            """
        )
        
        with gr.Tabs():
            with gr.Tab("📷 Image Upload"):
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### Input")
                        model_choice = gr.Dropdown(
                            choices=VALID_MODELS,
                            value=VALID_MODELS[0],
                            label="🔎 Select Model",
                            info="Choose a model for object detection or panoptic segmentation."
                        )
                        model_info = gr.Markdown(
                            f"**Model Info**: {MODEL_DESCRIPTIONS[VALID_MODELS[0]]}",
                            visible=True
                        )
                        image_input = gr.Image(type="pil", label="📷 Upload Image")
                        image_url_input = gr.Textbox(
                            label="🔗 Image URL",
                            placeholder="https://example.com/image.jpg"
                        )
                        with gr.Row():
                            submit_btn = gr.Button("✨ Detect", variant="primary")
                            clear_btn = gr.Button("🗑️ Clear", variant="secondary")
                        
                        model_choice.change(
                            fn=lambda model_name: f"**Model Info**: {MODEL_DESCRIPTIONS.get(model_name, 'No description available.')}",
                            inputs=model_choice,
                            outputs=model_info
                        )
                    
                    with gr.Column(scale=2):
                        gr.Markdown("### Results")
                        error_output = gr.Textbox(
                            label="⚠️ Errors",
                            visible=False,
                            lines=3,
                            max_lines=5
                        )
                        output_image = gr.Image(
                            type="pil",
                            label="🎯 Detected Image",
                            interactive=False
                        )
                        with gr.Row():
                            objects_output = gr.DataFrame(
                                label="📋 Detected Objects",
                                interactive=False,
                                value=None
                            )
                            unique_objects_output = gr.DataFrame(
                                label="🔍 Unique Objects",
                                interactive=False,
                                value=None
                            )
                        properties_output = gr.DataFrame(
                            label="📄 Image Properties",
                            interactive=False,
                            value=None
                        )
                
                def process_for_gradio(image, url, model_name):
                    try:
                        if image is None and not url:
                            return None, None, None, None, "Please provide an image or URL"
                        if image and url:
                            return None, None, None, None, "Please provide either an image or URL, not both"
                        
                        if url:
                            response = requests.get(url, timeout=10)
                            response.raise_for_status()
                            image = Image.open(BytesIO(response.content)).convert("RGB")
                        
                        detected_image, objects, scores, unique_objects, unique_scores, properties = process(image, model_name)
                        objects_df = pd.DataFrame({
                            "Object": objects,
                            "Confidence Score": [f"{score:.2f}" for score in scores]
                        }) if objects else pd.DataFrame(columns=["Object", "Confidence Score"])
                        unique_objects_df = pd.DataFrame({
                            "Unique Object": unique_objects,
                            "Confidence Score": [f"{score:.2f}" for score in unique_scores]
                        }) if unique_objects else pd.DataFrame(columns=["Unique Object", "Confidence Score"])
                        properties_df = pd.DataFrame([properties]) if properties else pd.DataFrame(columns=properties.keys())
                        return detected_image, objects_df, unique_objects_df, properties_df, ""
                    except Exception as e:
                        error_msg = f"Error processing image: {str(e)}"
                        logger.error(f"{error_msg}\n{traceback.format_exc()}")
                        return None, None, None, None, error_msg

                submit_btn.click(
                    fn=process_for_gradio,
                    inputs=[image_input, image_url_input, model_choice],
                    outputs=[output_image, objects_output, unique_objects_output, properties_output, error_output]
                )
                
                clear_btn.click(
                    fn=lambda: [None, "", None, None, None, None],
                    inputs=None,
                    outputs=[image_input, image_url_input, output_image, objects_output, unique_objects_output, properties_output, error_output]
                )
            
            with gr.Tab("🔗 URL Input"):
                gr.Markdown("### Process Image from URL")
                image_url_input = gr.Textbox(
                    label="🔗 Image URL",
                    placeholder="https://example.com/image.jpg"
                )
                url_model_choice = gr.Dropdown(
                    choices=VALID_MODELS,
                    value=VALID_MODELS[0],
                    label="🔎 Select Model"
                )
                url_model_info = gr.Markdown(
                    f"**Model Info**: {MODEL_DESCRIPTIONS[VALID_MODELS[0]]}",
                    visible=True
                )
                url_submit_btn = gr.Button("🔄 Process URL", variant="primary")
                url_output = gr.JSON(label="API Response")
                
                url_model_choice.change(
                    fn=lambda model_name: f"**Model Info**: {MODEL_DESCRIPTIONS.get(model_name, 'No description available.')}",
                    inputs=url_model_choice,
                    outputs=url_model_info
                )
                
                def process_url_for_gradio(url, model_name):
                    try:
                        response = requests.get(url, timeout=10)
                        response.raise_for_status()
                        image = Image.open(BytesIO(response.content)).convert("RGB")
                        detected_image, objects, scores, unique_objects, unique_scores, _ = process(image, model_name)
                        buffered = BytesIO()
                        detected_image.save(buffered, format="PNG")
                        img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
                        return {
                            "image_url": f"data:image/png;base64,{img_base64}",
                            "detected_objects": objects,
                            "confidence_scores": scores,
                            "unique_objects": unique_objects,
                            "unique_confidence_scores": unique_scores
                        }
                    except Exception as e:
                        error_msg = f"Error processing URL: {str(e)}"
                        logger.error(f"{error_msg}\n{traceback.format_exc()}")
                        return {"error": error_msg}

                url_submit_btn.click(
                    fn=process_url_for_gradio,
                    inputs=[image_url_input, url_model_choice],
                    outputs=[url_output]
                )
            
            with gr.Tab("ℹ️ Help"):
                gr.Markdown(
                    """
                    ## How to Use
                    - **Image Upload**: Select a model, upload an image or provide a URL, and click "Detect" to see detected objects and image properties.
                    - **URL Input**: Enter an image URL, select a model, and click "Process URL" to get results in JSON format.
                    - **Models**: Choose from DETR (object detection or panoptic segmentation) or YOLOS (lightweight detection).
                    - **Clear**: Reset all inputs and outputs using the "Clear" button.
                    - **Errors**: Check the error box for any processing issues.
                    
                    ## Tips
                    - Use high-quality images for better detection results.
                    - Panoptic models (e.g., DETR-ResNet-50-panoptic) provide segmentation masks for complex scenes.
                    - For faster processing, try YOLOS-Tiny on resource-constrained devices.
                    """
                )

    return demo

if __name__ == "__main__":
    demo = create_gradio_ui()
    demo.launch()
    # To run FastAPI, use: uvicorn object_detection:app --host 0.0.0.0 --port 8000