Added better bird classes and also the confidence score

app.py

@@ -61,14 +61,18 @@ def classify_image(image):
     classification = labels[predicted.item()]
 
     # Check if the predicted class is a bird
-    bird_classes = ['bird', 'fowl', 'hen', 'cock', 'rooster', 'peacock', 'parrot', 'eagle', 'owl', 'penguin']
-    is_bird = any(bird_class in classification.lower() for bird_class in bird_classes)
+    bird_categories = ['bird', 'fowl', 'hen', 'cock', 'rooster', 'peacock', 'parrot', 'eagle', 'owl', 'penguin']
+    is_bird = ('bird' in classification.lower()) or any(category in classification.lower() for category in bird_categories)
+
+    # Get the confidence score
+    confidence_score = torch.nn.functional.softmax(output[0], dim=0)[predicted].item()
+    confidence_percentage = f"{confidence_score:.2%}"
 
     if is_bird:
-        return f"This is a bird! Specifically, it looks like a {classification}."
+        return f"This is a bird! Specifically, it looks like a {classification}. Model confidence: {confidence_percentage}"
     else:
-        return f"This is not a bird. It appears to be a {classification}."
-
+        return f"This is not a bird. It appears to be a {classification}. Model confidence: {confidence_percentage}"
+#
 # Dynamically create the list of example images
 example_files = sorted(glob.glob("examples/*.png"))
 examples = [[file] for file in example_files]