|
|
import torch |
|
|
import gradio as gr |
|
|
from transformers import Owlv2Processor, Owlv2ForObjectDetection |
|
|
import spaces |
|
|
|
|
|
|
|
|
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() |
|
|
|
|
|
|
|
|
|
