app.py

import gradio as gr
import requests
import json
import os
import time
from collections import defaultdict
from PIL import Image
import io

BASE_URL = "https://api.jigsawstack.com/v1"
headers = {
    "x-api-key": os.getenv("JIGSAWSTACK_API_KEY")
}

# Rate limiting configuration
request_times = defaultdict(list)
MAX_REQUESTS = 20  # Maximum requests per time window
TIME_WINDOW = 3600   # Time window in seconds (1 hour)

def get_real_ip(request: gr.Request):
    """Extract real IP address using x-forwarded-for header or fallback"""
    if not request:
        return "unknown"
    
    forwarded = request.headers.get("x-forwarded-for")
    if forwarded:
        ip = forwarded.split(",")[0].strip()  # First IP in the list is the client's
    else:
        ip = request.client.host  # fallback
    return ip

def check_rate_limit(request: gr.Request):
    """Check if the current request exceeds rate limits"""
    if not request:
        return True, "Rate limit check failed - no request info"
    
    ip = get_real_ip(request)
    now = time.time()

    # Clean up old timestamps outside the time window
    request_times[ip] = [t for t in request_times[ip] if now - t < TIME_WINDOW]
    
    # Check if rate limit exceeded
    if len(request_times[ip]) >= MAX_REQUESTS:
        time_remaining = int(TIME_WINDOW - (now - request_times[ip][0]))
        time_remaining_minutes = round(time_remaining / 60, 1)      
        time_window_minutes = round(TIME_WINDOW / 60, 1)
        
        return False, f"Rate limit exceeded. You can make {MAX_REQUESTS} requests per {time_window_minutes} minutes. Try again in {time_remaining_minutes} minutes."
    
    # Add current request timestamp
    request_times[ip].append(now)
    return True, ""

def detect_objects(request: gr.Request, image_url=None, file_store_key=None, prompts=None, features=None):
    rate_limit_ok, rate_limit_msg = check_rate_limit(request)
    if not rate_limit_ok:
        return f"❌ {rate_limit_msg}", "", "", None
    
    if not image_url and not file_store_key:
        return "❌ Please provide either an image URL or file store key.", "", "", None

    if image_url and file_store_key:
        return "❌ Provide only one: image URL or file store key.", "", "", None

    try:
        payload = {}
        if image_url:
            payload["url"] = image_url
        if file_store_key:
            payload["file_store_key"] = file_store_key
        
        # Add optional parameters
        if prompts:
            payload["prompts"] = prompts
        if features:
            payload["features"] = features
        
        # Always return annotated image
        payload["annotated_image"] = True
        
        # Always use url as return_type
        payload["return_type"] = "url"

        response = requests.post(f"{BASE_URL}/ai/object_detection", headers=headers, json=payload)
        if response.status_code != 200:
            return f"❌ Error: {response.status_code} - {response.text}", "", "", None

        result = response.json()
        if not result.get("success"):
            return "❌ Detection failed.", "", "", None

        status = "✅ Detection successful!"
        objects = result.get("objects", [])
        annotated_image_url = result.get("annotated_image")
        
        # Create description with object details
        description = f"Image Size: {result.get('width', 'Unknown')} x {result.get('height', 'Unknown')}\n\n"
        description += f"Total Objects Detected: {len(objects)}\n\n"

        for i, obj in enumerate(objects):
            bounds = obj.get("bounds", {})
            label = obj.get("label", "Unknown")
            bound_text = ""
            
            if bounds:
                width = bounds.get("width", "Unknown")
                height = bounds.get("height", "Unknown")
                top_left = bounds.get("top_left", {})
                if top_left:
                    x, y = top_left.get("x", "?"), top_left.get("y", "?")
                    bound_text = f"Position: ({x}, {y}), Size: {width}x{height}"
            
            description += f"• {label}\n  {bound_text}\n"

        raw_json = json.dumps(result, indent=2)
        return status, description.strip(), raw_json, annotated_image_url

    except Exception as e:
        return f"❌ Error: {str(e)}", "", "", None

with gr.Blocks() as demo:
    gr.Markdown("""
    <div style='text-align: center; margin-bottom: 24px;'>
        <h1 style='font-size:2.2em; margin-bottom: 0.2em;'>🧩 Object Detection</h1>
        <p style='font-size:1.2em; margin-top: 0;'>Detect objects within images with great accuracy using AI models.</p>
        <p style='font-size:1em; margin-top: 0.5em;'>For more details and API usage, see the <a href='https://jigsawstack.com/docs/api-reference/ai/object-detection' target='_blank'>documentation</a>.</p>
    </div>
    """)

    with gr.Row():
        with gr.Column():
            input_type = gr.Radio(choices=["Image URL", "File Store Key"], value="Image URL", label="Input Type")
            image_url = gr.Textbox(label="Image URL", placeholder="https://example.com/image.jpg", visible=True)
            file_store_key = gr.Textbox(label="File Store Key", placeholder="my-image.jpg", visible=False)
            
            # Advanced options
            prompts = gr.Textbox(label="Prompts (comma-separated)", placeholder="wine glass, bottle, cup", info="Targeted object detection prompts")
            features = gr.CheckboxGroup(choices=["object_detection", "gui"], value=["object_detection"], label="Features")
            
            detect_btn = gr.Button("🔍 Detect Objects")
            clear_btn = gr.Button("Clear")

        with gr.Column():
            status_box = gr.Textbox(label="Status", interactive=False)
            desc_display = gr.Textbox(label="Object Details", lines=10, interactive=False)
            
            # Annotated image display - always visible
            annotated_image_display = gr.Image(label="Annotated Image")
            
            json_box = gr.Accordion("Raw JSON Response", open=False)
            with json_box:
                json_output = gr.Textbox(show_label=False, lines=20, interactive=False)

    def toggle_inputs(choice):
        return (
            gr.update(visible=(choice == "Image URL")),
            gr.update(visible=(choice == "File Store Key"))
        )

    input_type.change(fn=toggle_inputs, inputs=input_type, outputs=[image_url, file_store_key])

    def on_detect(input_mode, url, key, prompts_text, features_list, request: gr.Request):
        # Parse prompts
        prompts_list = None
        if prompts_text.strip():
            prompts_list = [p.strip() for p in prompts_text.split(",") if p.strip()]
        
        if input_mode == "Image URL":
            return detect_objects(
                request=request,
                image_url=url.strip(),
                prompts=prompts_list,
                features=features_list
            )
        else:
            return detect_objects(
                request=request,
                file_store_key=key.strip(),
                prompts=prompts_list,
                features=features_list
            )

    detect_btn.click(fn=on_detect, inputs=[
        input_type, image_url, file_store_key, prompts, features
    ], outputs=[status_box, desc_display, json_output, annotated_image_display])

    def clear_all():
        return "Image URL", "", "", "", "", ["object_detection"], "", "", "", None

    clear_btn.click(fn=clear_all, inputs=[], outputs=[
        input_type, image_url, file_store_key, prompts, features,
        status_box, desc_display, json_output, annotated_image_display
    ])

demo.launch()