Update gradcam.py

gradcam.py

@@ -74,9 +74,10 @@ class GradCam():
     def print_top_categories(self, model, img_tensor, top_k=5):
         logits = model(img_tensor.unsqueeze(0)).logits
         probabilities = torch.nn.functional.softmax(logits, dim=1)
-        indices = logits.cpu()[0, :].detach().numpy().argsort()[-top_k :][::-1]
-        for i in indices:
-            print(f"Predicted class (sorted from most confident) {i}: {model.config.id2label[i]}, confidence: {probabilities[0][i].item()}")
+        topIdx = logits.cpu()[0, :].detach().numpy().argsort()[-1]
+        topClass = model.config.id2label[topIdx]
+        topScore = probabilities[0][topIdx].item()
+        return {"label": topClass, "score": topScore}
 
     def reshape_transform_vit_huggingface(self, x):
         activations = x[:, 1:, :]