Click on an image to superimpose a heatmap

click_heatmap.py

import cv2
import numpy as np
import argparse

'''
usage: python click_heatmap.py <image file>
    left-click: add point to heatmap
    s: save image (000.png, 001.png, ...) 
    q: quit
    r: reset
'''

class Heatmapper:

    def __init__(self, img_path):
        self.img = cv2.imread(img_path)

        H, W = self.img.shape[0], self.img.shape[1]
        self.heatmap = np.zeros((H, W))
        
        k_size = int(np.sqrt(H*W)/5.0)
        if k_size%2 == 0:
            k_size += 1
        self.k_size = k_size

        self.overlay = None
        self.save_idx = 0

        cv2.namedWindow('img')
        cv2.imshow('img', self.img)
        cv2.setMouseCallback('img', self.collect_clicks)


    def collect_clicks(self, event, x, y, flags, param):

        if event == cv2.EVENT_LBUTTONDOWN:
            self.heatmap[y][x] += 1.0

            hmap = cv2.GaussianBlur(self.heatmap, (self.k_size, self.k_size), 0)
            hmap = hmap/hmap.max()
            hmap = (hmap*255).astype(np.uint8)
            hmap = cv2.applyColorMap(hmap, cv2.COLORMAP_JET)

            img = 0.7*self.img + 0.3*hmap
            img = img.astype(np.uint8)

            self.overlay = img
            cv2.imshow('img', img)


    def run(self):
        while True:
            key = cv2.waitKey(1) & 0xFF

            if ord('q')==key:
                return

            elif ord('s')==key:
                cv2.imwrite(f'{self.save_idx:03d}.png', self.overlay)
                print (f'Saved image to {self.save_idx:03d}.png')
                self.save_idx += 1

            elif ord('r')==key:
                self.heatmap.fill(0)
                cv2.imshow('img', self.img)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--img', required=True)
    args = parser.parse_args()

    Heatmapper(args.img).run()