Update gradcam.py

gradcam.py

@@ -84,55 +84,4 @@ class GradCam():
         activations = activations.view(activations.shape[0],
                                     14, 14, activations.shape[2])
         activations = activations.transpose(2, 3).transpose(1, 2)
-        return activations
-
-
-
-if __name__ == "__main__":
-
-    faceGrabber = FaceGrabber()
-    gradCam = GradCam()
-
-    image = Image.open("Feature-Image-74.jpg").convert("RGB")
-    face = faceGrabber.grab_faces(np.array(image))
-    if face is not None:
-        image = Image.fromarray(face)
-
-    img_tensor = transforms.ToTensor()(image)
-
-    model = ViTForImageClassification.from_pretrained("ongkn/attraction-classifier")
-    targets_for_gradcam = [ClassifierOutputTarget(gradCam.category_name_to_index(model, "pos")),
-                        ClassifierOutputTarget(gradCam.category_name_to_index(model, "neg"))]
-    target_layer_dff = model.vit.layernorm
-    target_layer_gradcam = model.vit.encoder.layer[-2].output
-    image_resized = image.resize((224, 224))
-    tensor_resized = transforms.ToTensor()(image_resized)
-
-    dff_image = run_dff_on_image(model=model,
-                                target_layer=target_layer_dff,
-                                classifier=model.classifier,
-                                img_pil=image_resized,
-                                img_tensor=tensor_resized,
-                                reshape_transform=gradCam.reshape_transform_vit_huggingface,
-                                n_components=5,
-                                top_k=10,
-                                threshold=0,
-                                output_size=None) #(500, 500))
-    cv.namedWindow("DFF Image", cv.WINDOW_KEEPRATIO)
-    cv.imshow("DFF Image", cv.cvtColor(dff_image, cv.COLOR_BGR2RGB))
-    cv.resizeWindow("DFF Image", 2500, 700)
-    # cv.waitKey(0)
-    # cv.destroyAllWindows()
-    grad_cam_image = gradCam.run_grad_cam_on_image(model=model,
-                                        target_layer=target_layer_gradcam,
-                                        targets_for_gradcam=targets_for_gradcam,
-                                        input_tensor=tensor_resized,
-                                        input_image=image_resized,
-                                        reshape_transform=gradCam.reshape_transform_vit_huggingface,
-                                        threshold=0)
-    cv.namedWindow("Grad-CAM Image", cv.WINDOW_KEEPRATIO)
-    cv.imshow("Grad-CAM Image", grad_cam_image)
-    cv.resizeWindow("Grad-CAM Image", 2000, 1250)
-    cv.waitKey(0)
-    cv.destroyAllWindows()
-    gradCam.print_top_categories(model, tensor_resized)
+        return activations