Update app.py

app.py

@@ -3,6 +3,7 @@ from hugsvision.inference.VisionClassifierInference import VisionClassifierInfer
 import gradio as gr
 import cv2
 import numpy as np
+from PIL import Image  # Importe a classe Image do módulo PIL (Pillow)
 
 # Load the pretrained ViT model and feature extractor
 path = "mrm8488/vit-base-patch16-224_finetuned-kvasirv2-colonoscopy"
@@ -15,18 +16,30 @@ classifier = VisionClassifierInference(
     model=model,
 )
 
-# Define a function to classify and overlay the label on the image
+# Define uma função para classificar e sobrepor o rótulo na imagem
 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
+    # Predict the label
+    label = classifier.predict(img_path=img)
+
+    # Load the image using OpenCV
+    image = cv2.imread(img)
+
+    # Add a white rectangle for the label
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    org = (10, 30)
+    font_scale = 1
+    color = (255, 255, 255)  # White color
+    thickness = 2
+    text_size = cv2.getTextSize(label, font, font_scale, thickness)[0]
+    cv2.rectangle(image, (org[0] - 10, org[1] - text_size[1] - 10), (org[0] + text_size[0], org[1]), color, cv2.FILLED)
+
+    # Put the label text on the white rectangle
+    cv2.putText(image, label, org, font, font_scale, (0, 0, 0), thickness, cv2.LINE_AA)
+
+    # Convert the image to RGB format for Gradio using PIL's Image class
+    image_pil = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+
+    return image_pil
 
 iface = gr.Interface(
     fn=classify_image_with_overlay,