import gradio as gr
from transformers import pipeline
from PIL import Image, ImageDraw

checkpoint = "google/owlvit-base-patch32"
detector = pipeline(model=checkpoint, task="zero-shot-object-detection")

def detect_and_visualize_objects(image):
    # Convert the image to RGB format
    image = image.convert("RGB")

    # Process the image using the object detection model
    predictions = detector(
        image,
        candidate_labels=["human face", "rocket"],
    )

    # Draw bounding boxes and labels on the image
    draw = ImageDraw.Draw(image)
    if len(predictions) == 0:
        draw.text((100, 100), "Object not found in image", fill="red")
    else:
        for prediction in predictions:
            box = prediction["box"]
            label = prediction["label"]
            score = prediction["score"]

            xmin, ymin, xmax, ymax = box.values()
            draw.rectangle((xmin, ymin, xmax, ymax), outline="red", width=1)
            draw.text((xmin, ymin), f"{label}: {round(score, 2)}", fill="white")

    # Return the annotated image
    return image

# Define the Gradio interface
image_input = gr.inputs.Image(type="pil")
image_output = gr.outputs.Image(type="pil")
iface = gr.Interface(
    fn=detect_and_visualize_objects,
    inputs=image_input,
    outputs=image_output,
    title="Space and War Missile Detection System",
    description="Detect objects in an image using a pre-trained model and visualize the results.",
    

)

# Launch the Gradio interface
iface.launch(debug=True)