两种方式构建heatmap

我会带着你远行 2022-10-05 05:58 153阅读 0赞

import numpy as np
    import cv2
    
    """
    固定尺寸
    """
    class HeatmapBuilder:
        def __init__(self, sigma):
            buffer_size = 100
            center = (buffer_size // 2, buffer_size // 2)
            x_range = [i for i in range(int(buffer_size))]
            y_range = [i for i in range(int(buffer_size))]
            xx, yy = np.meshgrid(x_range, y_range)
            d2 = (xx - center[0]) ** 2 + (yy - center[1]) ** 2
            exponent = d2 / 2.0 / sigma / sigma
            mask = exponent <= 4.6052
            mask_y, mask_x = np.where(mask)
    
            cofid_map = np.exp(-exponent)
            cofid_map[cofid_map > 1.0] = 1.0
            cofid_map = cofid_map.astype(np.float32)
            self.patch = cofid_map[np.min(mask_y):np.max(mask_y), np.min(mask_x):np.max(mask_x)]
    
        def build(self, img, label):
            # print(label)
            # print(label[0])
            # print(label[1])
            x = int(label[0])
            y = int(label[1])
    
            if x < 0 or y < 0 or x > img.shape[1] or y > img.shape[0]:
                return
    
            if x - self.patch.shape[1] // 2 > 0:
                x_start = 0
            else:
                x_start = self.patch.shape[1] // 2 - x
    
            if x + self.patch.shape[1] // 2 < img.shape[1]:
                x_end = self.patch.shape[1]
            else:
                x_end = img.shape[1] - x + self.patch.shape[1] // 2
    
            if y - self.patch.shape[0] // 2 > 0:
                y_start = 0
            else:
                y_start = self.patch.shape[0] // 2 - y
    
            if y + self.patch.shape[0] // 2 < img.shape[0]:
                y_end = self.patch.shape[0]
            else:
                y_end = img.shape[0] - y + self.patch.shape[0] // 2
    
            lt = (x - self.patch.shape[1] // 2, y - self.patch.shape[0] // 2)
            img[lt[1] + y_start:lt[1] + y_end, lt[0] + x_start:lt[0] + x_end] = \
                np.stack([img[lt[1] + y_start:lt[1] + y_end, lt[0] + x_start:lt[0] + x_end],
                          self.patch[y_start:y_end, x_start:x_end]], axis=-1).max(axis=-1)
    
            # cv2.imshow("heatmap",img)
            # cv2.waitKey(0)
            return img
    """
    可变大小
    """
    class HeatmapBuildernew:
        def __init__(self):
            pass
    
        def gaussian2D(self, shape, sigma=1):
            m, n = [(ss - 1.) / 2. for ss in shape]
            y, x = np.ogrid[-m:m+1,-n:n+1]
            h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
            h[h < np.finfo(h.dtype).eps * h.max()] = 0
            return h
    
        def draw_umich_gaussian(self, heatmap, center, radius, k=1):
            diameter = 2 * radius + 1
            gaussian = self.gaussian2D((diameter, diameter), sigma=diameter / 6)
            x, y = int(center[0]), int(center[1])
            height, width = heatmap.shape[0:2]
            left, right = min(x, radius), min(width - x, radius + 1)
            top, bottom = min(y, radius), min(height - y, radius + 1)
            masked_heatmap  = heatmap[y - top:y + bottom, x - left:x + right]
            masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right]
            if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0: # TODO debug
                np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)
            return heatmap
    
    black = np.zeros(shape=(64,64))
    
    H1 = HeatmapBuilder(1)
    # H2 = HeatmapBuildernew()
    
    img = H1.build(black,(32,32))
    # img = H2.draw_umich_gaussian(black,(32,32),10)
    cv2.imshow('',img)
    cv2.waitKey(0)