from transformers import ViTFeatureExtractor, ViTForImageClassification
from hugsvision.inference.VisionClassifierInference import VisionClassifierInference
import gradio as gr
import cv2
import numpy as np

# Load the pretrained ViT model and feature extractor
path = "mrm8488/vit-base-patch16-224_finetuned-kvasirv2-colonoscopy"
feature_extractor = ViTFeatureExtractor.from_pretrained(path)
model = ViTForImageClassification.from_pretrained(path)

# Create a VisionClassifierInference instance
classifier = VisionClassifierInference(
    feature_extractor=feature_extractor,
    model=model,
)

# Define a function to classify and overlay the label on the image
def classify_image_with_overlay(img):
    # Converte a imagem NumPy ndarray para um objeto Pillow Image
    img_pil = Image.fromarray(img)
    
    # Realiza a classificação usando o modelo e adiciona o rótulo previsto à imagem
    image_with_text = classificar_imagem(img_pil)
    
    # Converte a imagem resultante de volta para NumPy ndarray
    image_with_text_np = np.array(image_with_text)
    
    return image_with_text_np

iface = gr.Interface(
    fn=classify_image_with_overlay,
    inputs=gr.inputs.Image(),
    outputs=gr.outputs.Image(type="numpy"),
    live=True,
    title="ViT Image Classifier with Overlay",
    description="Upload an image for classification with label overlay.",
)


if __name__ == "__main__":
    iface.launch()