diff --git a/opennsfw2/_inspection.py b/opennsfw2/_inspection.py index 2fd7f17..2a67df3 100644 --- a/opennsfw2/_inspection.py +++ b/opennsfw2/_inspection.py @@ -71,7 +71,7 @@ def _resize( def save_grad_cam( - pil_image: Image.Image, + image: NDUInt8Array, heatmap: NDFloat32Array, grad_cam_path: str, alpha: float @@ -92,9 +92,9 @@ def save_grad_cam( jet_heatmap = jet_colors[scaled_heatmap] # Superimpose the heatmap on the input image after resizing. - jet_heatmap = _resize(jet_heatmap, pil_image.height, pil_image.width) + jet_heatmap = _resize(jet_heatmap, image.shape[0], image.shape[1]) - superimposed_image = jet_heatmap * alpha + np.array(pil_image) + superimposed_image = jet_heatmap * alpha + image pil_superimposed_image = array_to_img(superimposed_image) # Save the superimposed image. @@ -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 ) - save_grad_cam(pil_image, heatmap, grad_cam_path, alpha) + 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 index a05d02e..b266dd6 100644 Binary files a/tests/output_grad_cam_1.jpg and b/tests/output_grad_cam_1.jpg differ diff --git a/tests/output_grad_cam_2.jpg b/tests/output_grad_cam_2.jpg index 39f4698..f00d76b 100644 Binary files a/tests/output_grad_cam_2.jpg and b/tests/output_grad_cam_2.jpg differ diff --git a/tests/output_grad_cam_3.jpg b/tests/output_grad_cam_3.jpg index d136b35..da6fc1e 100644 Binary files a/tests/output_grad_cam_3.jpg and b/tests/output_grad_cam_3.jpg differ