Spaces:
Runtime error
Runtime error
File size: 2,348 Bytes
396610c |
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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import streamlit as st
import torch
from transformers import Owlv2Processor, Owlv2ForObjectDetection
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import random
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble").to(device)
processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemble")
st.title("Zero-Shot Object Detection with OWLv2")
uploaded_image = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
text_queries = st.text_input("Enter text queries (comma-separated):")
score_threshold = st.slider("Score Threshold", min_value=0.0, max_value=1.0, value=0.1, step=0.01)
def query_image(img, text_queries, score_threshold):
try:
img = Image.open(img).convert("RGB")
img_np = np.array(img)
text_queries = text_queries.split(",")
size = max(img_np.shape[:2])
target_sizes = torch.Tensor([[size, size]])
inputs = processor(text=text_queries, images=img_np, return_tensors="pt").to(device)
with torch.no_grad():
outputs = model(**inputs)
outputs.logits = outputs.logits.cpu()
outputs.pred_boxes = outputs.pred_boxes.cpu()
results = processor.post_process_object_detection(outputs=outputs, target_sizes=target_sizes)
boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["labels"]
result_labels = []
for box, score, label in zip(boxes, scores, labels):
box = [int(i) for i in box.tolist()]
if score < score_threshold:
continue
result_labels.append((box, text_queries[label.item()]))
return img, result_labels
except Exception as e:
st.error(f"Error performing object detection: {e}")
if uploaded_image is not None:
annotated_image, detected_objects = query_image(uploaded_image, text_queries, score_threshold)
if annotated_image:
draw = ImageDraw.Draw(annotated_image)
font = ImageFont.load_default()
for box, label in detected_objects:
color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
draw.rectangle(box, outline=color, width=3)
draw.text((box[0], box[1]), label, fill="black", font=font)
st.image(annotated_image, caption="Annotated Image", use_column_width=True)
|