File size: 2,147 Bytes
b63db34 c67e1ac b63db34 c67e1ac b63db34 c67e1ac |
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 |
import torch
import gradio as gr
from transformers import Owlv2Processor, Owlv2ForObjectDetection
import spaces
# Use GPU if available
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")
def query_image(Upload_Image, Text, score_threshold):
Text = Text
Text = Text.split(",")
size = max(Upload_Image.shape[:2])
target_sizes = torch.Tensor([[size, size]])
inputs = processor(text=Text, images=Upload_Image, 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[label.item()]))
return Upload_Image, result_labels
descricao = """
Você pode usar o AnyVision para consultar imagens com descrições de texto de qualquer objeto.
Para usá-lo, basta carregar uma imagem e inserir descrições de texto separadas por vírgula dos objetos que deseja consultar na imagem.
Você também pode usar o controle deslizante de limiar de pontuação para definir um limite para filtrar previsões de baixa probabilidade.
Você pode obter melhores previsões consultando a imagem com modelos de texto usados no treinamento do modelo original: por exemplo, *"foto de uma bandeira estrelada"*,
*"imagem de um sapato"*.
"""
demo = gr.Interface(
query_image,
inputs=[gr.Image(), "text", gr.Slider(0, 1, value=0.1)],
outputs="annotatedimage",
title="AnyVision - Detector de Objetos Zero-Shot com Owl2",
description=descricao
)
demo.launch()
|