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