import torch import cv2 import gradio as gr import numpy as np from transformers import OwlViTProcessor, OwlViTForObjectDetection # Use GPU if available if torch.cuda.is_available(): device = torch.device("cuda:4") else: device = torch.device("cpu") model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32").to(device) model.eval() processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32") def query_image(img, text_queries, score_threshold): text_queries = text_queries text_queries = text_queries.split(",") target_sizes = torch.Tensor([img.shape[:2]]) inputs = processor(text=text_queries, images=img, 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(outputs=outputs, target_sizes=target_sizes) boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["labels"] font = cv2.FONT_HERSHEY_SIMPLEX for box, score, label in zip(boxes, scores, labels): box = [int(i) for i in box.tolist()] if score >= score_threshold: img = cv2.rectangle(img, box[:2], box[2:], (255,0,0), 5) if box[3] + 25 > 768: y = box[3] - 10 else: y = box[3] + 25 img = cv2.putText( img, text_queries[label], (box[0], y), font, 1, (255,0,0), 2, cv2.LINE_AA ) return img description = """ Gradio demo for OWL-ViT, introduced in Simple Open-Vocabulary Object Detection with Vision Transformers. \n\nYou can use OWL-ViT to query images with text descriptions of any object. To use it, simply upload an image and enter comma separated text descriptions of objects you want to query the image for. You can also use the score threshold slider to set a threshold to filter out low probability predictions. \n\nOWL-ViT is trained on text templates, hence you can get better predictions by querying the image with text templates used in training the original model: *"photo of a star-spangled banner"*, *"image of a shoe"*. Refer to the CLIP paper to see the full list of text templates used to augment the training data. \n\nColab demo """ demo = gr.Interface( query_image, inputs=[gr.Image(), "text", gr.Slider(0, 1, value=0.1)], outputs="image", title="Zero-Shot Object Detection with OWL-ViT", description=description, examples=[ ["./demo_images/cats.png", "cats,ears", 0.11], ["./demo_images/demo1.jpg", "bear,soil,sea", 0.1], ["./demo_images/demo2.jpg", "dog,ear,leg,eyes,tail", 0.1], ["./demo_images/tanager.jpg", "wing,eyes,back,legs,tail", 0.01] ], ) # demo.launch() demo.launch(server_name="0.0.0.0", debug=True)