app.py

import argparse
import base64
import logging
import os
import sys
import threading
from collections import Counter
from io import BytesIO
from typing import Dict, List, Optional, Tuple, Union

import gradio as gr
import pandas as pd
import requests
import torch
import uvicorn
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import JSONResponse
from PIL import Image, ImageDraw, ImageStat
from transformers import (
    DetrForObjectDetection,
    DetrForSegmentation,
    DetrImageProcessor,
    YolosForObjectDetection,
    YolosImageProcessor,
)
import nest_asyncio

# ------------------------------
# Configuration
# ------------------------------

# Configure logging for debugging and monitoring
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

# Define constants for model and server configuration
CONFIDENCE_THRESHOLD: float = 0.5  # Default threshold for object detection confidence
VALID_MODELS: List[str] = [
    "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: Dict[str, str] = {
    "facebook/detr-resnet-50": "DETR with ResNet-50 for object detection. Fast and accurate.",
    "facebook/detr-resnet-101": "DETR with ResNet-101 for object detection. More accurate, slower.",
    "facebook/detr-resnet-50-panoptic": "DETR with ResNet-50 for panoptic segmentation.",
    "facebook/detr-resnet-101-panoptic": "DETR with ResNet-101 for panoptic segmentation.",
    "hustvl/yolos-tiny": "YOLOS Tiny. Lightweight and fast.",
    "hustvl/yolos-base": "YOLOS Base. Balances speed and accuracy."
}
DEFAULT_GRADIO_PORT: int = 7860  # Default port for Gradio UI
DEFAULT_FASTAPI_PORT: int = 8000  # Default port for FastAPI server
PORT_RANGE: range = range(7860, 7870)  # Range of ports to try for Gradio
MAX_PORT_ATTEMPTS: int = 10  # Maximum attempts to find an available port

# Thread-safe storage for lazy-loaded models and processors
models: Dict[str, any] = {}
processors: Dict[str, any] = {}
model_lock = threading.Lock()

# ------------------------------
# Image Processing
# ------------------------------

def process_image(
    image: Optional[Image.Image],
    url: Optional[str],
    model_name: str,
    for_json: bool = False,
    confidence_threshold: float = CONFIDENCE_THRESHOLD
) -> Union[Dict, Tuple[Optional[Image.Image], Optional[pd.DataFrame], Optional[pd.DataFrame], Optional[pd.DataFrame], str]]:
    """
    Process an image for object detection or panoptic segmentation, handling Gradio and FastAPI inputs.

    Args:
        image: PIL Image object from file upload (optional).
        url: URL of the image to process (optional).
        model_name: Name of the model to use (must be in VALID_MODELS).
        for_json: If True, return JSON dict for API/JSON tab; else, return tuple for Gradio Home tab.
        confidence_threshold: Minimum confidence score for detection (default: 0.5).

    Returns:
        For JSON: Dict with base64-encoded image, detected objects, and confidence scores.
        For Gradio: Tuple of (annotated image, objects DataFrame, unique objects DataFrame, properties DataFrame, error message).
    """
    try:
        # Validate input: ensure exactly one of image or URL is provided
        if image is None and not url:
            return {"error": "Please provide an image or URL"} if for_json else (None, None, None, None, "Please provide an image or URL")
        if image and url:
            return {"error": "Provide either an image or URL, not both"} if for_json else (None, None, None, None, "Provide either an image or URL, not both")
        if model_name not in VALID_MODELS:
            error_msg = f"Invalid model: {model_name}. Choose from: {VALID_MODELS}"
            return {"error": error_msg} if for_json else (None, None, None, None, error_msg)

        # Calculate margin threshold: (1 - confidence_threshold) / 2 + confidence_threshold
        margin_threshold = (1 - confidence_threshold) / 2 + confidence_threshold

        # Load image from URL if provided
        if url:
            response = requests.get(url, timeout=10)
            response.raise_for_status()
            image = Image.open(BytesIO(response.content)).convert("RGB")

        # Load model and processor thread-safely
        with model_lock:
            if model_name not in models:
                logger.info(f"Loading model: {model_name}")
                try:
                    # Select appropriate model and processor based on 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)
                except Exception as e:
                    error_msg = f"Failed to load model: {str(e)}"
                    logger.error(error_msg)
                    return {"error": error_msg} if for_json else (None, None, None, None, error_msg)
            model, processor = models[model_name], processors[model_name]

        # Prepare image for model processing
        inputs = processor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs)

        # Initialize drawing context for annotations
        draw = ImageDraw.Draw(image)
        object_names: List[str] = []
        confidence_scores: List[float] = []
        object_counter = Counter()
        target_sizes = torch.tensor([image.size[::-1]])

        # Process results based on model type (panoptic or object detection)
        if "panoptic" in model_name:
            # Handle panoptic segmentation
            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)
                # Apply segmentation mask if available
                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:
            # Handle object detection
            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()
                    label_name = model.config.id2label.get(label.item(), "Unknown")
                    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]
                    # Use yellow for confidence_threshold <= score < margin_threshold, green for >= margin_threshold
                    color = "#FFFF00" if score < margin_threshold else "#32CD32"
                    draw.rectangle([x, y, x2, y2], outline=color, width=2)
                    draw.text((x2 - text_width - 2, y - text_height - 2), text, fill=color)
                    object_names.append(label_name)
                    confidence_scores.append(float(score))
                    object_counter[label_name] = float(score)

        # Compile unique objects and their highest confidence scores
        unique_objects = list(object_counter.keys())
        unique_confidences = [object_counter[obj] for obj in unique_objects]

        # Calculate image properties (metadata)
        properties: Dict[str, str] = {
            "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": "Unknown",
            "Mean (R,G,B)": "Unknown",
            "StdDev (R,G,B)": "Unknown",
        }
        try:
            # Compute file size
            buffered = BytesIO()
            image.save(buffered, format="PNG")
            properties["File Size"] = f"{len(buffered.getvalue()) / 1024:.2f} KB"
            # Compute color statistics
            stat = ImageStat.Stat(image)
            properties["Mean (R,G,B)"] = ", ".join(f"{m:.2f}" for m in stat.mean)
            properties["StdDev (R,G,B)"] = ", ".join(f"{s:.2f}" for s in stat.stddev)
        except Exception as e:
            logger.error(f"Error calculating image stats: {str(e)}")

        # Prepare output based on request type
        if for_json:
            # Return JSON with base64-encoded image
            buffered = BytesIO()
            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": object_names,
                "confidence_scores": confidence_scores,
                "unique_objects": unique_objects,
                "unique_confidence_scores": unique_confidences,
            }
        else:
            # Return tuple for Gradio Home tab with DataFrames
            objects_df = (
                pd.DataFrame({"Object": object_names, "Confidence Score": [f"{score:.2f}" for score in confidence_scores]})
                if object_names 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_confidences]})
                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 image, objects_df, unique_objects_df, properties_df, ""

    except requests.RequestException as e:
        # Handle URL fetch errors
        error_msg = f"Error fetching image from URL: {str(e)}"
        logger.error(f"{error_msg}\n{traceback.format_exc()}")
        return {"error": error_msg} if for_json else (None, None, None, None, error_msg)
    except Exception as e:
        # Handle general processing errors
        error_msg = f"Error processing image: {str(e)}"
        logger.error(f"{error_msg}\n{traceback.format_exc()}")
        return {"error": error_msg} if for_json else (None, None, None, None, error_msg)

# ------------------------------
# FastAPI Setup
# ------------------------------

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

@app.post("/detect")
async def detect_objects_endpoint(
    file: Optional[UploadFile] = File(None),
    image_url: Optional[str] = Form(None),
    model_name: str = Form(VALID_MODELS[0]),
    confidence_threshold: float = Form(CONFIDENCE_THRESHOLD),
) -> JSONResponse:
    """
    FastAPI endpoint to detect objects in an image from file upload or URL.

    Args:
        file: Uploaded image file (optional).
        image_url: URL of the image (optional).
        model_name: Model to use for detection (default: first VALID_MODELS entry).
        confidence_threshold: Confidence threshold for detection (default: 0.5).

    Returns:
        JSONResponse with base64-encoded image, detected objects, and confidence scores.

    Raises:
        HTTPException: For invalid inputs or processing errors.
    """
    try:
        # Validate input: ensure exactly one of file or URL
        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, not both.")
        # Validate confidence threshold
        if not 0 <= confidence_threshold <= 1:
            raise HTTPException(status_code=400, detail="Confidence threshold must be between 0 and 1.")
        # Load image from file if provided
        image = None
        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")
        # Process image with specified parameters
        result = process_image(image, image_url, model_name, for_json=True, confidence_threshold=confidence_threshold)
        if "error" in result:
            raise HTTPException(status_code=400, detail=result["error"])
        return JSONResponse(content=result)
    except HTTPException:
        raise
    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 Setup
# ------------------------------

def create_gradio_ui() -> gr.Blocks:
    """
    Create and configure the Gradio UI for object detection with Home, JSON, and Help tabs.

    Returns:
        Gradio Blocks object representing the UI.

    Raises:
        RuntimeError: If UI creation fails.
    """
    try:
        # Initialize Gradio Blocks with a custom theme
        with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="gray")) as demo:
            # Display app header
            gr.Markdown(
                f"""
                # 🚀 Object Detection App
                Upload an image or provide a URL to detect objects using transformer models (DETR, YOLOS).
                Running on port: {os.getenv('GRADIO_SERVER_PORT', 'auto-selected')}
                """
            )

            # Create tabbed interface
            with gr.Tabs():
                # Home tab (formerly Image Upload)
                with gr.Tab("🏠 Home"):
                    with gr.Row():
                        # Left column for inputs
                        with gr.Column(scale=1):
                            gr.Markdown("### Input")
                            # Model selection dropdown
                            model_choice = gr.Dropdown(choices=VALID_MODELS, value=VALID_MODELS[0], label="🔎 Select Model")
                            model_info = gr.Markdown(f"**Model Info**: {MODEL_DESCRIPTIONS[VALID_MODELS[0]]}")
                            # Image upload input
                            image_input = gr.Image(type="pil", label="📷 Upload Image")
                            # Image URL input
                            image_url_input = gr.Textbox(label="🔗 Image URL", placeholder="https://example.com/image.jpg")
                            # Buttons for submission and clearing
                            with gr.Row():
                                submit_btn = gr.Button("✨ Detect", variant="primary")
                                clear_btn = gr.Button("🗑️ Clear", variant="secondary")

                            # Update model info when model changes
                            model_choice.change(
                                fn=lambda model_name: f"**Model Info**: {MODEL_DESCRIPTIONS.get(model_name, 'No description available.')}",
                                inputs=model_choice,
                                outputs=model_info,
                            )

                        # Right column for results
                        with gr.Column(scale=2):
                            gr.Markdown("### Results")
                            # Error display (hidden by default)
                            error_output = gr.Textbox(label="⚠️ Errors", visible=False, lines=3, max_lines=5)
                            # Annotated image output
                            output_image = gr.Image(type="pil", label="🎯 Detected Image", interactive=False)
                            # Detected and unique objects tables
                            with gr.Row():
                                objects_output = gr.DataFrame(label="📋 Detected Objects", interactive=False)
                                unique_objects_output = gr.DataFrame(label="🔍 Unique Objects", interactive=False)
                            # Image properties table
                            properties_output = gr.DataFrame(label="📄 Image Properties", interactive=False)

                    # Process image when Detect button is clicked
                    submit_btn.click(
                        fn=process_image,
                        inputs=[image_input, image_url_input, model_choice],
                        outputs=[output_image, objects_output, unique_objects_output, properties_output, error_output],
                    )

                    # Clear all inputs and outputs
                    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],
                    )

                # JSON tab for API-like output
                with gr.Tab("🔗 JSON"):
                    with gr.Row():
                        # Left column for inputs
                        with gr.Column(scale=1):
                            gr.Markdown("### Process Image for JSON")
                            # Model selection dropdown
                            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]]}")
                            # Image upload input
                            image_input_json = gr.Image(type="pil", label="📷 Upload Image")
                            # Image URL input
                            image_url_input_json = gr.Textbox(label="🔗 Image URL", placeholder="https://example.com/image.jpg")
                            # Process button
                            url_submit_btn = gr.Button("🔄 Process", variant="primary")

                            # Update model info when model changes
                            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,
                            )

                        # Right column for JSON output
                        with gr.Column(scale=1):
                            # JSON output display
                            url_output = gr.JSON(label="API Response")

                    # Process image and return JSON when Process button is clicked
                    url_submit_btn.click(
                        fn=lambda img, url, model: process_image(img, url, model, for_json=True),
                        inputs=[image_input_json, image_url_input_json, url_model_choice],
                        outputs=[url_output],
                    )

                # Help tab with usage instructions
                with gr.Tab("ℹ️ Help"):
                    gr.Markdown(
                        """
                        ## How to Use
                        - **Home**: Select a model, upload an image or provide a URL, click "Detect" to see results.
                        - **JSON**: Select a model, upload an image or enter a URL, click "Process" for JSON output.
                        - **Models**: Choose DETR (detection or panoptic) or YOLOS (lightweight detection).
                        - **Clear**: Reset inputs/outputs in Home tab.
                        - **Errors**: Check error box (Home) or JSON response (JSON) for issues.

                        ## Tips
                        - Use high-quality images for better results.
                        - Panoptic models provide segmentation masks for complex scenes.
                        - YOLOS-Tiny is faster for resource-constrained devices.
                        """
                    )

        return demo

    except Exception as e:
        logger.error(f"Error creating Gradio UI: {str(e)}\n{traceback.format_exc()}")
        raise RuntimeError(f"Failed to create Gradio UI: {str(e)}")

# ------------------------------
# Launcher
# ------------------------------

def parse_args() -> argparse.Namespace:
    """
    Parse command-line arguments for configuring the application.

    Returns:
        Parsed arguments as a Namespace object.
    """
    parser = argparse.ArgumentParser(description="Object Detection App with Gradio and FastAPI.")
    # Gradio port argument
    parser.add_argument("--gradio-port", type=int, default=DEFAULT_GRADIO_PORT, help=f"Gradio port (default: {DEFAULT_GRADIO_PORT}).")
    # FastAPI enable flag
    parser.add_argument("--enable-fastapi", action="store_true", help="Enable FastAPI server.")
    # FastAPI port argument
    parser.add_argument("--fastapi-port", type=int, default=DEFAULT_FASTAPI_PORT, help=f"FastAPI port (default: {DEFAULT_FASTAPI_PORT}).")
    # Confidence threshold argument
    parser.add_argument("--confidence-threshold", type=float, default=CONFIDENCE_THRESHOLD, help="Confidence threshold for detection (default: 0.5).")
    # Parse known arguments, ignoring unrecognized ones
    args, _ = parser.parse_known_args()
    # Validate confidence threshold
    if not 0 <= args.confidence_threshold <= 1:
        parser.error("Confidence threshold must be between 0 and 1.")
    return args

def find_available_port(start_port: int, port_range: range, max_attempts: int) -> Optional[int]:
    """
    Find an available port within the specified range.

    Args:
        start_port: Initial port to try.
        port_range: Range of ports to attempt.
        max_attempts: Maximum number of ports to try.

    Returns:
        Available port number, or None if no port is found.
    """
    import socket
    # Check environment variable for port override
    port = int(os.getenv("GRADIO_SERVER_PORT", start_port))
    attempts = 0
    while attempts < max_attempts:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            try:
                # Attempt to bind to the port
                s.bind(("0.0.0.0", port))
                logger.debug(f"Port {port} is available")
                return port
            except OSError as e:
                if e.errno == 98:  # Port in use
                    logger.debug(f"Port {port} is in use")
                    port = port + 1 if port < max(port_range) else min(port_range)
                    attempts += 1
                else:
                    raise
    logger.error(f"No available port in range {min(port_range)}-{max(port_range)}")
    return None

def main() -> None:
    """
    Launch the Gradio UI and optional FastAPI server.

    Raises:
        SystemExit: On interruption or critical errors.
    """
    try:
        # Apply nest_asyncio for compatibility with Jupyter/Colab
        nest_asyncio.apply()
        # Parse command-line arguments
        args = parse_args()
        logger.info(f"Parsed arguments: {args}")
        # Find available port for Gradio
        gradio_port = find_available_port(args.gradio_port, PORT_RANGE, MAX_PORT_ATTEMPTS)
        if gradio_port is None:
            logger.error("Failed to find an available port for Gradio UI")
            sys.exit(1)

        # Start FastAPI server in a thread if enabled
        if args.enable_fastapi:
            logger.info(f"Starting FastAPI on port {args.fastapi_port}")
            fastapi_thread = threading.Thread(
                target=lambda: uvicorn.run(app, host="0.0.0.0", port=args.fastapi_port),
                daemon=True
            )
            fastapi_thread.start()

        # Launch Gradio UI
        logger.info(f"Starting Gradio UI on port {gradio_port}")
        demo = create_gradio_ui()
        demo.launch(server_port=gradio_port, server_name="0.0.0.0")

    except KeyboardInterrupt:
        logger.info("Application terminated by user.")
        sys.exit(0)
    except Exception as e:
        logger.error(f"Error: {str(e)}\n{traceback.format_exc()}")
        sys.exit(1)

if __name__ == "__main__":
    main()