Update app.py

app.py

@@ -1,13 +1,6 @@
 import gradio as gr
-import torch
 from transformers import ViTFeatureExtractor, ViTForImageClassification
-import cv2
-import numpy as np 
-
-# Carregue o modelo ViT
-model_name = "mrm8488/vit-base-patch16-224_finetuned-kvasirv2-colonoscopy"
-feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
-model = ViTForImageClassification.from_pretrained(model_name)
+import numpy as np
 
 # Mapeamento de classe ID para rótulo
 id2label = {
@@ -21,52 +14,30 @@ id2label = {
     "7": "ulcerative-colitis"
 }
 
+# Carregue o modelo ViT
+model_name = "mrm8488/vit-base-patch16-224_finetuned-kvasirv2-colonoscopy"
+feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
+model = ViTForImageClassification.from_pretrained(model_name)
+
 # Função para classificar a imagem
 def classify_image(input_image):
-    # Redimensionar a imagem de entrada para ser 2x maior
-    input_image = cv2.resize(input_image, None, fx=2, fy=2)
-    
     # Pré-processar a imagem usando o extrator de características
     inputs = feature_extractor(input_image, return_tensors="pt")
     # Realizar inferência com o modelo
     outputs = model(**inputs)
     # Obter a classe prevista
-    predicted_class_id = torch.argmax(outputs.logits, dim=1).item()
-    # Obter a probabilidade da classe prevista
-    predicted_class_prob = torch.softmax(outputs.logits, dim=1)[0, predicted_class_id].item()
-    # Converter o ID da classe em rótulo usando o mapeamento id2label
+    predicted_class_id = np.argmax(outputs.logits)
+    # Obter o rótulo da classe a partir do mapeamento id2label
     predicted_class_label = id2label.get(str(predicted_class_id), "Desconhecido")
-    
-    # Converter a imagem de numpy para BGR (formato OpenCV)
-    input_image_bgr = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
-    
-    # Definir cores de borda para cada classe (aqui, cores aleatórias)
-    class_colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0),
-                    (255, 0, 255), (0, 255, 255), (128, 128, 128), (0, 0, 0)]
-    
-    # Adicionar uma borda colorida à imagem
-    border_color = class_colors[predicted_class_id]
-    input_image_bgr = cv2.copyMakeBorder(input_image_bgr, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=border_color)
-    
-    # Adicionar o rótulo da previsão na imagem
-    font = cv2.FONT_HERSHEY_SIMPLEX
-    text = f'Classe: {predicted_class_label} ({predicted_class_prob:.2f})'
-    text_size = cv2.getTextSize(text, font, 0.7, 2)[0]
-    text_x = (input_image_bgr.shape[1] - text_size[0]) // 2
-    text_y = input_image_bgr.shape[0] - 30  # Ajuste da posição vertical
-    cv2.putText(input_image_bgr, text, (text_x, text_y), font, 0.7, (255, 255, 255), 2, cv2.LINE_AA)
-    
-    # Converter a imagem resultante de volta para RGB (formato Pillow)
-    result_image = cv2.cvtColor(input_image_bgr, cv2.COLOR_BGR2RGB)
-    return result_image
+    return predicted_class_label
 
 # Criar uma interface Gradio
 interface = gr.Interface(
     fn=classify_image,
     inputs=gr.inputs.Image(type="numpy", label="Carregar uma imagem"),
-    outputs=gr.outputs.Image(type="numpy", label="Resultado"),
+    outputs=gr.outputs.Label(num_top_classes=1),
     title="Classificador de Imagem ViT",
-    description="Esta aplicação Gradio permite classificar imagens usando um modelo Vision Transformer (ViT). O rótulo da previsão está na imagem de saída."
+    description="Esta aplicação Gradio permite classificar imagens usando um modelo Vision Transformer (ViT).",
 )
 
 # Iniciar a aplicação Gradio