Update app.py

app.py

@@ -1,8 +1,7 @@
 import gradio as gr
 import torch
 from transformers import ViTFeatureExtractor, ViTForImageClassification
-from PIL import Image, ImageDraw, ImageFont
-import io
+import cv2
 import numpy as np 
 
 # Carregue o modelo ViT
@@ -33,27 +32,19 @@ def classify_image(input_image):
     # Converter o ID da classe em rótulo usando o mapeamento id2label
     predicted_class_label = id2label.get(str(predicted_class_id), "Desconhecido")
     
-    # Abrir a imagem usando PIL
-    image = Image.fromarray(input_image.astype('uint8'))
+    # Converter a imagem de numpy para BGR (formato OpenCV)
+    input_image_bgr = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
     
-    # Obter as dimensões da imagem
-    width, height = image.size
-    
-    # Criar uma imagem com a faixa branca na parte inferior
-    result_image = Image.new('RGB', (width, height + 40), color='white')
-    result_image.paste(image, (0, 0))
-    
-    # Adicionar o rótulo da previsão na faixa branca
-    draw = ImageDraw.Draw(result_image)
-    font = ImageFont.load_default()
+    # Adicionar o rótulo da previsão na imagem
+    font = cv2.FONT_HERSHEY_SIMPLEX
     text = f'Previsão: {predicted_class_label}'
-    text_width, text_height = draw.textsize(text, font=font)
-    x = (width - text_width) // 2
-    y = height + 10
-    draw.text((x, y), text, fill='black', font=font)
+    text_size = cv2.getTextSize(text, font, 0.5, 1)[0]
+    text_x = (input_image_bgr.shape[1] - text_size[0]) // 2
+    text_y = input_image_bgr.shape[0] - 10
+    cv2.putText(input_image_bgr, text, (text_x, text_y), font, 0.5, (255, 255, 255), 1)
     
-    # Converter a imagem resultante de volta para numpy
-    result_image = np.array(result_image)
+    # Converter a imagem resultante de volta para RGB (formato Pillow)
+    result_image = cv2.cvtColor(input_image_bgr, cv2.COLOR_BGR2RGB)
     return result_image
 
 # Criar uma interface Gradio
@@ -62,7 +53,7 @@ interface = gr.Interface(
     inputs=gr.inputs.Image(type="numpy", label="Carregar uma imagem"),
     outputs=gr.outputs.Image(type="numpy", label="Resultado"),
     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 parte inferior da imagem de saída em uma faixa branca."
+    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."
 )
 
 # Iniciar a aplicação Gradio