Code optimization

app.py

@@ -27,7 +27,8 @@ test_loader = dataset.get_test_data_loader(**dataloader_args)
 classes = ('plane', 'car', 'bird', 'cat', 'deer',
            'dog', 'frog', 'horse', 'ship', 'truck')
 
-cache_dict = {"misclassified_images": None, "is_misclassified_images": None, "num_misclassified_images": None}
+# Get the misclassified data from test dataset
+misclassified_data = get_misclassified_data(model, device, test_loader)
 
 def resize_image_pil(image, new_width, new_height):
 
@@ -86,20 +87,11 @@ def inference(input_img, is_grad_cam=True, transparency = 0.5, target_layer_numb
     # Pick the top n predictions
     top_n_confidences = dict(list(sorted_confidences.items())[:top_predictions])
 
-    if (is_misclassified_images != cache_dict["is_misclassified_images"] or
-         num_misclassified_images != cache_dict["num_misclassified_images"]):
-        cache_dict["is_misclassified_images"] = is_misclassified_images
-        cache_dict["num_misclassified_images"] = num_misclassified_images
-        if is_misclassified_images:
-            # Get the misclassified data from test dataset
-            misclassified_data = get_misclassified_data(model, device, test_loader)
-            # Plot the misclassified data
-            misclassified_images = display_cifar_misclassified_data(misclassified_data, number_of_samples=num_misclassified_images)
-            cache_dict["misclassified_images"] = misclassified_images
-        else:
-            misclassified_images = None
+    if is_misclassified_images:
+        # Plot the misclassified data
+        misclassified_images = display_cifar_misclassified_data(misclassified_data, number_of_samples=num_misclassified_images)
     else:
-        misclassified_images = cache_dict["misclassified_images"]
+        misclassified_images = None
     
     return classes[prediction[0].item()], visualization, top_n_confidences, misclassified_images
 

utils.py

@@ -51,6 +51,9 @@ def get_misclassified_data(model, device, test_loader):
     with torch.no_grad():
         # Extract images, labels in a batch
         for data, target in test_loader:
+            
+            if len(misclassified_data) > 40:
+                break
 
             # Migrate the data to the device
             data, target = data.to(device), target.to(device)