Update app.py

app.py

@@ -1,30 +1,26 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+from hugsvision.inference.VisionClassifierInference import VisionClassifierInference
 
-# Carregue a pipeline de classificação de imagem ViT
-image_classifier = pipeline("image-classification", model="mrm8488/vit-base-patch16-224_finetuned-kvasirv2-colonoscopy")
+# Load the pre-trained ViT model
+path = "mrm8488/vit-base-patch16-224_finetuned-kvasirv2-colonoscopy"
+classifier = VisionClassifierInference(
+    feature_extractor=ViTFeatureExtractor.from_pretrained(path),
+    model=ViTForImageClassification.from_pretrained(path),
+)
 
-# Função para classificar a imagem
+# Define a Gradio interface
 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
+    label = classifier.predict(img_path=img)
+    return f"Predicted class: {label}"
 
-# Interface Gradio
 iface = gr.Interface(
     fn=classify_image,
-    inputs=gr.inputs.Image(),
-    outputs=gr.outputs.Textbox(),
-    live=True,
-    title="ViT Image Classifier",
-    description="Carregue uma imagem para classificação.",
+    inputs=gr.inputs.Image(type="file", label="Upload an image"),
+    outputs="text",
+    title="Image Classifier",
+    description="Classify images using a pre-trained ViT model",
 )
 
-if __name__ == "__main__":
-    iface.launch()
+# Launch the Gradio app
+iface.launch()