Update app.py

app.py

@@ -3,6 +3,7 @@ from torchvision import models, transforms
 from PIL import Image
 import gradio as gr
 
+# Updated class names with 'plaque' in front of 'calculus' and 'gingivitis'
 class_names = [
     "plaque_calculus",
     "caries",
@@ -33,18 +34,37 @@ def predict_image(image):
 
     with torch.no_grad():
         outputs = model(processed_image)
-        _, top_indices = torch.topk(outputs, 2)  # Get top 2 predictions
-        top_classes = [class_names[idx] for idx in top_indices[0]]
+        probabilities = torch.nn.functional.softmax(outputs, dim=1)
+        top_probs, top_indices = torch.topk(probabilities, 2)  # Get top 2 predictions
 
-    return ", ".join(top_classes)
+        top_class_1 = class_names[top_indices[0][0]]
+        top_prob_1 = top_probs[0][0].item()
+
+        # Initialize result with the top prediction
+        result = top_class_1
+
+        # Include the second prediction if the top prediction's probability is less than 80%
+        if top_prob_1 < 0.8:
+            top_class_2 = class_names[top_indices[0][1]]
+            result += f", {top_class_2}"
+
+    return result
+
+# Example images to use as input
+examples = [
+    ["example_image1.jpg"],
+    ["example_image2.jpg"],
+    ["example_image3.jpg"]
+]
 
 # Set up the Gradio interface
 iface = gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil"),
-    outputs="text",  # Output will be text listing 
-    title="Dental Classifier",
-    description="Upload an image to predict its class."
+    outputs="text",  # Output will be text listing the predictions
+    title="Medical Image Classification",
+    description="Upload an image to predict its class. Displays the top 2 predictions if needed.",
+    examples=examples  # Add example images
 )
 
 # Launch the interface