import gradio as gr
import os
import json
import google.generativeai as genai
from PIL import Image, ImageDraw, ImageFont, ImageColor
from dotenv import load_dotenv

load_dotenv()
api_key = os.getenv("Gemini_API_Key") 
genai.configure(api_key=api_key)

SYSTEM_INSTRUCTION = """
Return bounding boxes as a JSON array with labels. Never return masks or code fencing. Limit to 25 objects.
If an object is present multiple times, name them according to their unique characteristic.
Use the format: [{"box_2d": [ymin, xmin, ymax, xmax], "label": "object_name"}]
"""

model = genai.GenerativeModel(
    model_name='gemini-2.5-flash', 
    system_instruction=SYSTEM_INSTRUCTION
)

def parse_response(text):
    """Cleans markdown formatting if the model includes it."""
    clean = text.replace("json", "").replace("", "").strip()
    return clean

def draw_boxes(image, boxes_json):
    """Draws boxes on the image using normalized coordinates (0-1000)."""
    img = image.copy()
    draw = ImageDraw.Draw(img)
    width, height = img.size
    colors = ["red", "green", "blue", "yellow", "magenta", "cyan", "orange"]

    try:
        data = json.loads(boxes_json)
        for i, item in enumerate(data):
            color = colors[i % len(colors)]
            # Gemini order: [ymin, xmin, ymax, xmax]
            ymin, xmin, ymax, xmax = item["box_2d"]
            
            left = (xmin / 1000) * width
            top = (ymin / 1000) * height
            right = (xmax / 1000) * width
            bottom = (ymax / 1000) * height
            
            draw.rectangle([left, top, right, bottom], outline=color, width=5)
            draw.text((left + 5, top + 5), item.get("label", "Object"), fill=color)
    except Exception as e:
        print(f"Error drawing boxes: {e}")
    
    return img

def process_image(input_img, user_instruction):
    if input_img is None:
        return None
    
    if not user_instruction.strip():
        user_instruction = "Identify and label the objects in the image."

    # Combine user instruction with JSON formatting requirement
    full_prompt = f"{user_instruction} Return ONLY the JSON array of bounding boxes."
    
    response = model.generate_content([full_prompt, input_img])
    json_data = parse_response(response.text)
    
    output_img = draw_boxes(input_img, json_data)
    return output_img

demo = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Textbox(
            label="What should I find?", 
            placeholder="e.g., 'Find all the red apples' or 'Detect the cars and pedestrians'",
            lines=2
        )
    ],
    outputs=gr.Image(type="pil", label="Image with Bounding Boxes"),
    title="Your Own Object Detector",
    description="Upload an image and give specific instructions on what to detect."
)

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