Update app.py

app.py

@@ -2,29 +2,43 @@ import gradio as gr
 import torch
 from transformers import ViTFeatureExtractor, ViTForImageClassification
 
-# Load the ViT model
+# Carregue o modelo ViT
 model_name = "google/vit-base-patch16-224-in21k"
 feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
 model = ViTForImageClassification.from_pretrained(model_name)
 
-# Define a function to preprocess the image and classify it
+# Mapeamento de classe ID para rótulo
+id2label = {
+    "0": "dyed-lifted-polyps",
+    "1": "dyed-resection-margins",
+    "2": "esophagitis",
+    "3": "normal-cecum",
+    "4": "normal-pylorus",
+    "5": "normal-z-line",
+    "6": "polyps",
+    "7": "ulcerative-colitis"
+}
+
+# Função para classificar a imagem
 def classify_image(input_image):
-    # Preprocess the image using the feature extractor
+    # Pré-processar a imagem usando o extrator de características
     inputs = feature_extractor(input_image, return_tensors="pt")
-    # Perform inference with the model
+    # Realizar inferência com o modelo
     outputs = model(**inputs)
-    # Get the predicted label
-    predicted_class = torch.argmax(outputs.logits, dim=1).item()
-    return predicted_class
+    # Obter a classe prevista
+    predicted_class_id = torch.argmax(outputs.logits, dim=1).item()
+    # Converter o ID da classe em rótulo usando o mapeamento id2label
+    predicted_class_label = id2label.get(str(predicted_class_id), "Desconhecido")
+    return predicted_class_label
 
-# Create a Gradio interface
+# Criar uma interface Gradio
 interface = gr.Interface(
     fn=classify_image,
-    inputs=gr.inputs.Image(type="numpy", label="Upload an image"),
+    inputs=gr.inputs.Image(type="numpy", label="Carregar uma imagem"),
     outputs="label",
-    title="ViT Image Classifier",
-    description="This Gradio app allows you to classify images using a Vision Transformer (ViT) model."
+    title="Classificador de Imagem ViT",
+    description="Esta aplicação Gradio permite classificar imagens usando um modelo Vision Transformer (ViT)."
 )
 
-# Launch the Gradio app
+# Iniciar a aplicação Gradio
 interface.launch()