Spaces:
Paused
Paused
File size: 1,698 Bytes
4b3b02a |
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 |
from transformers import DetrImageProcessor, DetrForObjectDetection
from io import BytesIO
from PIL import Image
import requests
import torch
class Generador():
def __init__(self, configuraciones):
self.modelo = configuraciones.get('model')
self.tokenizer = configuraciones.get('tokenizer')
def generar_prediccion(self, imagen_bytes):
# @title **Ejemplo práctico**
prediccion = str()
try:
# Inicializamos los procesadores y el modelo
procesador = DetrImageProcessor.from_pretrained(self.tokenizer)
modelo = DetrForObjectDetection.from_pretrained(self.modelo)
# Procesamos nuestra imagen y objetos
inputs = procesador(images=imagen_bytes, return_tensors="pt")
outputs = modelo(**inputs)
objetos = torch.tensor([imagen_bytes.size[::-1]])
# filtramos objetos con probabilidad mayor al 90%
resultados = procesador.post_process_object_detection(
outputs, target_sizes=objetos, threshold=0.9)[0]
# imprimimos cada resultado con su probabilidad
for probabilidad, etiqueta, ubicacion in zip(resultados["scores"], resultados["labels"], resultados["boxes"]):
ubicacion = [round(i, 2) for i in ubicacion.tolist()]
_text = f"Encontrado:\t {modelo.config.id2label[etiqueta.item()]}\nProbabilidad: {round(probabilidad.item(), 3)} en ubicación: {ubicacion}."
prediccion = prediccion + _text
except Exception as error:
print(f"No es Chems\n{error}")
prediccion = str(error)
finally:
self.prediccion = prediccion
|