Updated app.py with confidence threshold of 0.8

app.py

@@ -6,6 +6,8 @@ import torch
 import torchvision.models as models
 from torch import nn
 
+from model import *
+
 from albumentations import (
     HorizontalFlip, VerticalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,
     Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,
@@ -98,11 +100,21 @@ def predict_image(image):
     # Assuming the output is a tensor representing class probabilities
     probabilities = torch.nn.functional.softmax(output[0], dim=0).numpy()
 
+    conf_threshold = 0.8
+
     # Get the class with the highest probability
     predicted_class = np.argmax(probabilities)
 
+    # If the probability is less than the threshold, return unknown else return the class
+    if probabilities[predicted_class] < conf_threshold:
+        string_to_return = "Predicted Class: Unknown"
+
+    else:
+        # Return the class label
+        string_to_return = f"Predicted Class: {classes[predicted_class+1]} with probability: {probabilities[predicted_class] * 100:.2f}%"
+
     # Return the class label
-    return "Predicted Class: " + classes[predicted_class+1]
+    return string_to_return
 
 # create a gradio interface
 gr.Interface(fn=predict_image, inputs="image", outputs="text").launch()