Update app.py

app.py

@@ -1,45 +1,48 @@
 import gradio as gr
 from fastai.vision.all import *
 
-cat_breeds = [
-    'Abyssinian', 'Bengal', 'Birman', 'Bombay', 'British_Shorthair',
-    'Egyptian_Mau', 'Main_Coon', 'Persian', 'Ragdoll', 'Russian_Blue',
-    'Siamese', 'Sphynx'
-]
+def is_cat(x): return x[0].isupper()
 
 def get_labels(path):
-    regex_pattern = r'(.+)_(\d+)\.jpg$'
-    match = re.match(regex_pattern, str(path.name))
-    if match:
-        label = match.group(1)
-        pet_type = 'cat' if label in cat_breeds else 'dog'
-        return [pet_type, label]
-    else:
-        # Handle the case where the regex doesn't match
-        print(f"Could not match pattern in {path.name}")
-        return []
-
-learn = load_learner('dog_cat_multi.pkl')
+    pet_type = 'cat' if is_cat(path.name) else 'dog'
+    breed = RegexLabeller(r'(.+)_\d+.jpg$')(path.name)
+
+    return [pet_type, breed]
+
+learn = load_learner('dog_cat_multi_v3.pkl')
 
 def predict(image):
     # Transform the image
     pil_image = PILImage.create(image)
-    
+
     # Predict
-    preds, _, probs = learn.predict(pil_image)
-    
+    preds, mask, probs = learn.predict(pil_image)
+
     # Apply the threshold
-    threshold = 0.425
-    classes = [learn.dls.vocab[i] for i in range(len(probs)) if probs[i] > threshold]
+    threshold = 0.570
+    classes_with_probs = [(learn.dls.vocab[i], probs[i].item()) for i in range(len(mask)) if probs[i] > threshold]
+
+    # Check if the prediction includes "dog" or "cat"
+    pet_type = None
+    breed = 'unknown' # Default to 'unknown' if no breed is identified
+    pet_prob = 0
+    breed_prob = 0
+    for class_name, prob in classes_with_probs:
+        if class_name == 'dog' or class_name == 'cat':
+            pet_type = class_name
+            pet_prob = prob
+        else:
+            breed = class_name
+            breed_prob = prob
+
+    # Check if pet_type is None (i.e., neither dog nor cat)
+    if pet_type is None:
+        return "This is not a cat, nor a dog."
 
-    pet_type = "dog" if "dog" in classes else "cat" if "cat" in classes else None
-    breeds = [breed for breed in classes if breed not in ["cat", "dog"]]
+    result = f"This is a {pet_type}, the breed is {breed if breed else 'unknown'}.\n" \
+             f"The probability for it being a {pet_type} is {pet_prob * 100:.2f}%, the probability of being the breed is {breed_prob * 100:.2f}%."
 
-    if pet_type:
-        breed_info = f"The breed is {'/'.join(breeds)}." if breeds else "The breed is not identified."
-        return f"This is a {pet_type}. {breed_info}"
-    else:
-        return "This is not a cat, nor a dog."
+    return result
 
 
 # Define the Gradio interface
@@ -50,7 +53,7 @@ iface = gr.Interface(
     live=True,
     title="Cat and Dog Image Classifier",
     description="Upload an image of a cat or a dog, and the model will identify the type and breed.",
-    article="This model has been trained on the Oxford Pets dataset and might not recognize all types dog and cat breeds. For best results, use clear images."
+    article="This model has been trained on the Oxford Pets dataset and might not recognize all types dog and cat breeds. For best results, use clear images of pets."
 )