Spaces:
Runtime error
Runtime error
File size: 1,525 Bytes
7ac295d 558764d 7ac295d 3573af2 beebfc8 3573af2 7ac295d 3573af2 e81fde2 70c390b 4fbbf16 7ac295d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
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)
|