Update visualization

visualization

@@ -33,14 +33,14 @@ model = vit_LRP(pretrained=True).cuda()
 model.eval()
 attribution_generator = LRP(model)
 
-def generate_visualization(original_image, class_index=None):
+def generate_visualization(original_image, class_index=None, use_threshold=False):
     transformer_attribution = attribution_generator.generate_LRP(original_image.unsqueeze(0).cuda(), method="transformer_attribution", index=class_index).detach()
     transformer_attribution = transformer_attribution.reshape(1, 1, 14, 14)
     transformer_attribution = torch.nn.functional.interpolate(transformer_attribution, scale_factor=16, mode='bilinear')
     transformer_attribution = transformer_attribution.reshape(224, 224).data.cpu().numpy()
     transformer_attribution = (transformer_attribution - transformer_attribution.min()) / (transformer_attribution.max() - transformer_attribution.min())
 
-    if use_thresholding:
+    if use_threshold:
       transformer_attribution = transformer_attribution * 255
       transformer_attribution = transformer_attribution.astype(np.uint8)
       ret, transformer_attribution = cv2.threshold(transformer_attribution, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)