Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from transformers import ViTFeatureExtractor, ViTForImageClassification
 from hugsvision.inference.VisionClassifierInference import VisionClassifierInference
+from PIL import Image, ImageDraw, ImageFont
 
 # Load the pre-trained ViT model
 path = "mrm8488/vit-base-patch16-224_finetuned-kvasirv2-colonoscopy"
@@ -9,7 +10,7 @@ classifier = VisionClassifierInference(
     model=ViTForImageClassification.from_pretrained(path),
 )
 
-# Define a Gradio interface
+
 def classify_image(image_file):
     """Classify an image using a pre-trained ViT model."""
     label = classifier.predict(img_path=image_file.name)
@@ -17,9 +18,19 @@ def classify_image(image_file):
     confidence = classifier.predict_proba(img_path=image_file.name)[0][label]
 
     # Get the PIL Image object for the uploaded image
-    image = image_file.read()
+    image = Image.open(image_file)
+
+    # Draw the predicted label on the image
+    draw = ImageDraw.Draw(image)
+    font = ImageFont.truetype("arial.ttf", 20)
+    draw.text((10, 10), f"Predicted class: {label} (confidence: {confidence:.2f})", font=font, fill=(255, 255, 255))
+
+    # Save the modified image to a BytesIO object
+    output_image = BytesIO()
+    image.save(output_image, format="JPEG")
+    output_image.seek(0)
 
-    return image, f"Predicted class: {label} (confidence: {confidence:.2f})"
+    return output_image, f"Predicted class: {label} (confidence: {confidence:.2f})"
 
 iface = gr.Interface(
     fn=classify_image,