set module guard

gradcam.py

@@ -81,48 +81,54 @@ class GradCam():
         activations = activations.transpose(2, 3).transpose(1, 2)
         return activations
 
-faceGrabber = FaceGrabber()
-
-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(category_name_to_index(model, "pos")),
-                       ClassifierOutputTarget(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=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 = 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=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()
-print_top_categories(model, tensor_resized)
+
+
+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 = 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)
+