limit predictions

app.py

@@ -16,30 +16,35 @@ model.load_state_dict(torch.load("best_model.pth", map_location=torch.device('cp
 model.eval()
 
 def predict(input_image):
-    #apply the transform
-    image_tensor = val_test_transform(input_image)
+    # 1. Transform the image (resize, normalize, etc.)
+    processed_image = val_test_transform(input_image)
 
-    #add batch dimension
-    image_tensor = image_tensor.unsqueeze(0)
+    # 2. Add a batch dimension because the model expects [batch, channels, height, width]
+    processed_image = processed_image.unsqueeze(0)
 
-    #run inference 
+    # 3. Run the image through the model
     with torch.no_grad():
-        output = model(image_tensor)
-    
-    #get probabilitiees
-    probabilities = torch.nn.functional.softmax(output,dim=1)[0]
+        model_output = model(processed_image)
+
+    # 4. Convert raw model scores into probabilities
+    probabilities = torch.nn.functional.softmax(model_output, dim=1)[0]
+
+    # 5. Choose how many results you want to show
+    number_of_predictions_to_show = 5
 
-    #create the output dictionary
-# Convert to dictionary
-    all_results = {CLASS_NAMES[i]: probabilities[i].item() for i in range(len(probabilities))}
+    # 6. Get the top-k highest probability classes
+    top_probabilities, top_class_indices = torch.topk(probabilities, number_of_predictions_to_show)
 
-# Sort by probability (highest first)
-    sorted_results = dict(sorted(all_results.items(), key=lambda x: x[1], reverse=True))
+    # 7. Build a dictionary of the results
+    results = {}
+    for rank in range(number_of_predictions_to_show):
+        class_index = top_class_indices[rank].item()
+        class_name = CLASS_NAMES[class_index]
+        probability_value = top_probabilities[rank].item()
 
-# Keep only top 10
-    top_10 = dict(list(sorted_results.items())[:10])
+        results[class_name] = probability_value
 
-    return top_10
+    return results
 
 
 with gr.Blocks(title="Plant Disease Classifier") as app: