Update app.py

app.py

@@ -1,61 +1,29 @@
-from transformers import ViTFeatureExtractor, ViTForImageClassification
-from hugsvision.inference.VisionClassifierInference import VisionClassifierInference
 import gradio as gr
-import cv2
-import numpy as np
-from PIL import Image
+from transformers import pipeline
 
-# 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)
+# Carregue a pipeline de classificação de imagem ViT
+image_classifier = pipeline("image-classification", model="mrm8488/vit-base-patch16-224_finetuned-kvasirv2-colonoscopy")
 
-# 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):
-    # Convert the numpy array image to a PIL image
-    img_pil = Image.fromarray(img)
-
-    # Predict the label and probability
-    prediction = classifier.predict_image(img=img_pil)
-
-    # Load the image using OpenCV
-    image = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
-
-    # Add a white rectangle for the label
-    font = cv2.FONT_HERSHEY_SIMPLEX
-    org = (10, 50)  # Adjust the position of the label
-    font_scale = 1.5  # Increase the font size
-    color = (255, 255, 255)  # White color
-    thickness = 2
-    text = f"{prediction['label']}: {prediction['probability']:.2f}"
-
-    text_size = cv2.getTextSize(text, 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, text, org, font, font_scale, (0, 0, 0), thickness, cv2.LINE_AA)
-
-    # Resize the image to a larger size
-    image = cv2.resize(image, (800, 800))  # Defina o tamanho desejado aqui
-
-    return image
-
-# ...
+# Função para classificar a imagem
+def classify_image(img):
+    # Execute a classificação da imagem usando a pipeline
+    result = image_classifier(img)
+    label = result[0]['label']
+    probability = result[0]['score']
 
+    # Formate o resultado para exibição
+    output_text = f"Classe: {label}, Probabilidade: {probability:.2f}"
+    
+    return output_text
 
+# Interface Gradio
 iface = gr.Interface(
-    fn=classify_image_with_overlay,
+    fn=classify_image,
     inputs=gr.inputs.Image(),
-    outputs=gr.outputs.Image(type="numpy"),
+    outputs=gr.outputs.Textbox(),
     live=True,
-    title="ViT Image Classifier with Overlay",
-    description="Upload an image for classification with label overlay.",
+    title="ViT Image Classifier",
+    description="Carregue uma imagem para classificação.",
 )
 
 if __name__ == "__main__":