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")
articles ="The Space and War Missile Detection System is a network of sensors that is used to detect objects in space and in the atmosphere. The system includes ground-based radars, satellites, and other sensors. The sensors are used to detect the heat, light, and radio waves emitted by objects in space. The system can also detect the movement of objects in space.",
 
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.",
    article=article,

)

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