generate_heatmap.py

import PIL.Image as Image
import cv2
import numpy as np


def main(data, image, n_bins, sigma_percentage):
    height, width, *_ = ref_image.shape
    if width > height:
        n_bins_x = int(n_bins)
        n_bins_y = int(n_bins * height / width)
    else:
        n_bins_x = int(n_bins * width / height)
        n_bins_y = int(n_bins)

    sigma = n_bins * sigma_percentage / 100

    if len(data):
        x, y = data.T
        y = 1 - y
        histogram, *_ = np.histogram2d(x, y, [n_bins_x, n_bins_y], [[0, 1], [0, 1]])
        histogram = histogram.T
        histogram = cv2.GaussianBlur(histogram, (0, 0), sigmaX=sigma, sigmaY=sigma)

        histogram_max = histogram.max()
        histogram *= (255.0 / histogram_max) if histogram_max else 0.0
        histogram = histogram.astype(np.uint8)
    else:
        histogram = np.zeros((n_bins_y, n_bins_x), dtype=np.uint8)

    heatmap = cv2.applyColorMap(histogram, cv2.COLORMAP_JET)
    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)

    heatmap_RGBA = np.ones((n_bins_y, n_bins_x, 4), dtype=np.uint8)
    heatmap_RGBA[:, :, 0:3] = heatmap
    heatmap_RGBA[:, :, 3] = histogram
    heatmap_RGBA = cv2.resize(
        heatmap_RGBA, dsize=(width, height), interpolation=cv2.INTER_LINEAR
    )

    heatmap_PIL = Image.fromarray(heatmap_RGBA, "RGBA")
    overlay_PIL = Image.fromarray(image, "RGB")
    overlay_PIL.paste(heatmap_PIL, (0, 0), heatmap_PIL)
    overlay_PIL.save("overlay.png")
    heatmap_PIL.save("heatmap.png")

    return heatmap_RGBA[:, :, 0:3]