Fix the border of the Grad-CAM images with YAHOO preprocessing

bhky · Aug 4, 2024 · 8ade275 · 8ade275
1 parent 7fd90b8
commit 8ade275
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 0 deletions.
diff --git a/opennsfw2/_inspection.py b/opennsfw2/_inspection.py
@@ -112,4 +112,12 @@ def make_and_save_nsfw_grad_cam(
         preprocess_image(pil_image, preprocessing), open_nsfw_model,
         "activation_stage3_block2", "fc_nsfw", 1
     )
+    if preprocessing == Preprocessing.YAHOO:
+        border_reduction_factor = 224 / 256
+        w, h = pil_image.size
+        w_new, h_new = int(w * border_reduction_factor), int(h * border_reduction_factor)
+        left_new = (w - w_new) // 2
+        upper_new = (h - h_new) // 2
+        pil_image = pil_image.crop((left_new, upper_new, left_new + w_new, upper_new + h_new))
+
     save_grad_cam(np.array(pil_image), heatmap, grad_cam_path, alpha)
diff --git a/tests/output_grad_cam_1.jpg b/tests/output_grad_cam_1.jpg
diff --git a/tests/output_grad_cam_2.jpg b/tests/output_grad_cam_2.jpg
diff --git a/tests/output_grad_cam_3.jpg b/tests/output_grad_cam_3.jpg