mmpose/codecs/utils/post_processing.py

# Copyright (c) OpenMMLab. All rights reserved.
from itertools import product
from typing import Tuple

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from torch import Tensor

from scipy.signal import convolve2d


def get_simcc_normalized(batch_pred_simcc, sigma=None):
    """Normalize the predicted SimCC.

    Args:
        batch_pred_simcc (torch.Tensor): The predicted SimCC.
        sigma (float): The sigma of the Gaussian distribution.

    Returns:
        torch.Tensor: The normalized SimCC.
    """
    B, K, _ = batch_pred_simcc.shape

    # Scale and clamp the tensor
    if sigma is not None:
        batch_pred_simcc = batch_pred_simcc / (sigma * np.sqrt(np.pi * 2))
    batch_pred_simcc = batch_pred_simcc.clamp(min=0)

    # Compute the binary mask
    mask = (batch_pred_simcc.amax(dim=-1) > 1).reshape(B, K, 1)

    # Normalize the tensor using the maximum value
    norm = (batch_pred_simcc / batch_pred_simcc.amax(dim=-1).reshape(B, K, 1))

    # Apply normalization
    batch_pred_simcc = torch.where(mask, norm, batch_pred_simcc)

    return batch_pred_simcc


def get_simcc_maximum(simcc_x: np.ndarray,
                      simcc_y: np.ndarray,
                      apply_softmax: bool = False
                      ) -> Tuple[np.ndarray, np.ndarray]:
    """Get maximum response location and value from simcc representations.

    Note:
        instance number: N
        num_keypoints: K
        heatmap height: H
        heatmap width: W

    Args:
        simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx)
        simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy)
        apply_softmax (bool): whether to apply softmax on the heatmap.
            Defaults to False.

    Returns:
        tuple:
        - locs (np.ndarray): locations of maximum heatmap responses in shape
            (K, 2) or (N, K, 2)
        - vals (np.ndarray): values of maximum heatmap responses in shape
            (K,) or (N, K)
    """

    assert isinstance(simcc_x, np.ndarray), ('simcc_x should be numpy.ndarray')
    assert isinstance(simcc_y, np.ndarray), ('simcc_y should be numpy.ndarray')
    assert simcc_x.ndim == 2 or simcc_x.ndim == 3, (
        f'Invalid shape {simcc_x.shape}')
    assert simcc_y.ndim == 2 or simcc_y.ndim == 3, (
        f'Invalid shape {simcc_y.shape}')
    assert simcc_x.ndim == simcc_y.ndim, (
        f'{simcc_x.shape} != {simcc_y.shape}')

    if simcc_x.ndim == 3:
        N, K, Wx = simcc_x.shape
        simcc_x = simcc_x.reshape(N * K, -1)
        simcc_y = simcc_y.reshape(N * K, -1)
    else:
        N = None

    if apply_softmax:
        simcc_x = simcc_x - np.max(simcc_x, axis=1, keepdims=True)
        simcc_y = simcc_y - np.max(simcc_y, axis=1, keepdims=True)
        ex, ey = np.exp(simcc_x), np.exp(simcc_y)
        simcc_x = ex / np.sum(ex, axis=1, keepdims=True)
        simcc_y = ey / np.sum(ey, axis=1, keepdims=True)

    x_locs = np.argmax(simcc_x, axis=1)
    y_locs = np.argmax(simcc_y, axis=1)
    locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
    max_val_x = np.amax(simcc_x, axis=1)
    max_val_y = np.amax(simcc_y, axis=1)

    mask = max_val_x > max_val_y
    max_val_x[mask] = max_val_y[mask]
    vals = max_val_x
    locs[vals <= 0.] = -1

    if N:
        locs = locs.reshape(N, K, 2)
        vals = vals.reshape(N, K)

    return locs, vals


def get_heatmap_3d_maximum(heatmaps: np.ndarray
                           ) -> Tuple[np.ndarray, np.ndarray]:
    """Get maximum response location and value from heatmaps.

    Note:
        batch_size: B
        num_keypoints: K
        heatmap dimension: D
        heatmap height: H
        heatmap width: W

    Args:
        heatmaps (np.ndarray): Heatmaps in shape (K, D, H, W) or
            (B, K, D, H, W)

    Returns:
        tuple:
        - locs (np.ndarray): locations of maximum heatmap responses in shape
            (K, 3) or (B, K, 3)
        - vals (np.ndarray): values of maximum heatmap responses in shape
            (K,) or (B, K)
    """
    assert isinstance(heatmaps,
                      np.ndarray), ('heatmaps should be numpy.ndarray')
    assert heatmaps.ndim == 4 or heatmaps.ndim == 5, (
        f'Invalid shape {heatmaps.shape}')

    if heatmaps.ndim == 4:
        K, D, H, W = heatmaps.shape
        B = None
        heatmaps_flatten = heatmaps.reshape(K, -1)
    else:
        B, K, D, H, W = heatmaps.shape
        heatmaps_flatten = heatmaps.reshape(B * K, -1)

    z_locs, y_locs, x_locs = np.unravel_index(
        np.argmax(heatmaps_flatten, axis=1), shape=(D, H, W))
    locs = np.stack((x_locs, y_locs, z_locs), axis=-1).astype(np.float32)
    vals = np.amax(heatmaps_flatten, axis=1)
    locs[vals <= 0.] = -1

    if B:
        locs = locs.reshape(B, K, 3)
        vals = vals.reshape(B, K)

    return locs, vals


def get_heatmap_maximum(heatmaps: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """Get maximum response location and value from heatmaps.

    Note:
        batch_size: B
        num_keypoints: K
        heatmap height: H
        heatmap width: W

    Args:
        heatmaps (np.ndarray): Heatmaps in shape (K, H, W) or (B, K, H, W)

    Returns:
        tuple:
        - locs (np.ndarray): locations of maximum heatmap responses in shape
            (K, 2) or (B, K, 2)
        - vals (np.ndarray): values of maximum heatmap responses in shape
            (K,) or (B, K)
    """
    assert isinstance(heatmaps,
                      np.ndarray), ('heatmaps should be numpy.ndarray')
    assert heatmaps.ndim == 3 or heatmaps.ndim == 4, (
        f'Invalid shape {heatmaps.shape}')

    if heatmaps.ndim == 3:
        K, H, W = heatmaps.shape
        B = None
        heatmaps_flatten = heatmaps.reshape(K, -1)
    else:
        B, K, H, W = heatmaps.shape
        heatmaps_flatten = heatmaps.reshape(B * K, -1)

    y_locs, x_locs = np.unravel_index(
        np.argmax(heatmaps_flatten, axis=1), shape=(H, W))
    locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
    vals = np.amax(heatmaps_flatten, axis=1)
    locs[vals <= 0.] = -1

    if B:
        locs = locs.reshape(B, K, 2)
        vals = vals.reshape(B, K)

    return locs, vals


def gaussian_blur(heatmaps: np.ndarray, kernel: int = 11) -> np.ndarray:
    """Modulate heatmap distribution with Gaussian.

    Note:
        - num_keypoints: K
        - heatmap height: H
        - heatmap width: W

    Args:
        heatmaps (np.ndarray[K, H, W]): model predicted heatmaps.
        kernel (int): Gaussian kernel size (K) for modulation, which should
            match the heatmap gaussian sigma when training.
            K=17 for sigma=3 and k=11 for sigma=2.

    Returns:
        np.ndarray ([K, H, W]): Modulated heatmap distribution.
    """
    assert kernel % 2 == 1

    border = (kernel - 1) // 2
    K, H, W = heatmaps.shape

    for k in range(K):
        origin_max = np.max(heatmaps[k])
        dr = np.zeros((H + 2 * border, W + 2 * border), dtype=np.float32)
        dr[border:-border, border:-border] = heatmaps[k].copy()
        dr = cv2.GaussianBlur(dr, (kernel, kernel), 0)
        heatmaps[k] = dr[border:-border, border:-border].copy()
        heatmaps[k] *= origin_max / (np.max(heatmaps[k])+1e-12)
    return heatmaps


def gaussian_blur1d(simcc: np.ndarray, kernel: int = 11) -> np.ndarray:
    """Modulate simcc distribution with Gaussian.

    Note:
        - num_keypoints: K
        - simcc length: Wx

    Args:
        simcc (np.ndarray[K, Wx]): model predicted simcc.
        kernel (int): Gaussian kernel size (K) for modulation, which should
            match the simcc gaussian sigma when training.
            K=17 for sigma=3 and k=11 for sigma=2.

    Returns:
        np.ndarray ([K, Wx]): Modulated simcc distribution.
    """
    assert kernel % 2 == 1

    border = (kernel - 1) // 2
    N, K, Wx = simcc.shape

    for n, k in product(range(N), range(K)):
        origin_max = np.max(simcc[n, k])
        dr = np.zeros((1, Wx + 2 * border), dtype=np.float32)
        dr[0, border:-border] = simcc[n, k].copy()
        dr = cv2.GaussianBlur(dr, (kernel, 1), 0)
        simcc[n, k] = dr[0, border:-border].copy()
        simcc[n, k] *= origin_max / np.max(simcc[n, k])
    return simcc


def batch_heatmap_nms(batch_heatmaps: Tensor, kernel_size: int = 5):
    """Apply NMS on a batch of heatmaps.

    Args:
        batch_heatmaps (Tensor): batch heatmaps in shape (B, K, H, W)
        kernel_size (int): The kernel size of the NMS which should be
            a odd integer. Defaults to 5

    Returns:
        Tensor: The batch heatmaps after NMS.
    """

    assert isinstance(kernel_size, int) and kernel_size % 2 == 1, \
        f'The kernel_size should be an odd integer, got {kernel_size}'

    padding = (kernel_size - 1) // 2

    maximum = F.max_pool2d(
        batch_heatmaps, kernel_size, stride=1, padding=padding)
    maximum_indicator = torch.eq(batch_heatmaps, maximum)
    batch_heatmaps = batch_heatmaps * maximum_indicator.float()

    return batch_heatmaps


def get_heatmap_expected_value(heatmaps: np.ndarray, parzen_size: float = 0.1, return_heatmap: bool = False) -> Tuple[np.ndarray, np.ndarray]:
    """Get maximum response location and value from heatmaps.

    Note:
        batch_size: B
        num_keypoints: K
        heatmap height: H
        heatmap width: W

    Args:
        heatmaps (np.ndarray): Heatmaps in shape (K, H, W) or (B, K, H, W)

    Returns:
        tuple:
        - locs (np.ndarray): locations of maximum heatmap responses in shape
            (K, 2) or (B, K, 2)
        - vals (np.ndarray): values of maximum heatmap responses in shape
            (K,) or (B, K)
    """
    assert isinstance(heatmaps,
                      np.ndarray), ('heatmaps should be numpy.ndarray')
    assert heatmaps.ndim == 3 or heatmaps.ndim == 4, (
        f'Invalid shape {heatmaps.shape}')
    
    assert parzen_size >= 0.0 and parzen_size <= 1.0, (
        f'Invalid parzen_size {parzen_size}')

    if heatmaps.ndim == 3:
        K, H, W = heatmaps.shape
        B = 1
        FIRST_DIM = K
        heatmaps_flatten = heatmaps.reshape(1, K, H, W)
    else:
        B, K, H, W = heatmaps.shape
        FIRST_DIM = K*B
        heatmaps_flatten = heatmaps.reshape(B, K, H, W)

    # Blur heatmaps with Gaussian
    # heatmaps_flatten = gaussian_blur(heatmaps_flatten, kernel=9)
    
    # Zero out pixels far from the maximum for each heatmap
    # heatmaps_tmp = heatmaps_flatten.copy().reshape(B*K, H*W)
    # y_locs, x_locs = np.unravel_index(
    #     np.argmax(heatmaps_tmp, axis=1), shape=(H, W))
    # locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
    # heatmaps_flatten = heatmaps_flatten.reshape(B*K, H, W)
    # for i, x in enumerate(x_locs):
    #     y = y_locs[i]
    #     start_x = int(max(0, x - 0.2*W))
    #     end_x = int(min(W, x + 0.2*W))
    #     start_y = int(max(0, y - 0.2*H))
    #     end_y = int(min(H, y + 0.2*H))
    #     mask = np.zeros((H, W))
    #     mask[start_y:end_y, start_x:end_x] = 1
    #     heatmaps_flatten[i] = heatmaps_flatten[i] * mask
    # heatmaps_flatten = heatmaps_flatten.reshape(B, K, H, W)


    bbox_area = np.sqrt(H/1.25 * W/1.25)

    kpt_sigmas = np.array(
        [2.6, 2.5, 2.5, 3.5, 3.5, 7.9, 7.9, 7.2, 7.2, 6.2, 6.2, 10.7, 10.7, 8.7, 8.7, 8.9, 8.9])/100
    
    heatmaps_covolved = np.zeros_like(heatmaps_flatten)
    for k in range(K):
        vars = (kpt_sigmas[k]*2)**2
        s = vars * bbox_area * 2
        s = np.clip(s, 0.55, 3.0)
        radius = np.ceil(s * 3).astype(int)
        diameter = 2*radius + 1
        diameter = np.ceil(diameter).astype(int)
        # kernel_sizes[kernel_sizes % 2 == 0] += 1
        center = diameter // 2
        dist_x = np.arange(diameter) - center
        dist_y = np.arange(diameter) - center
        dist_x, dist_y = np.meshgrid(dist_x, dist_y)
        dist = np.sqrt(dist_x**2 + dist_y**2)
        oks_kernel = np.exp(-dist**2 / (2 * s))
        oks_kernel = oks_kernel / oks_kernel.sum()
        
        htm = heatmaps_flatten[:, k, :, :].reshape(-1, H, W)
        # htm = np.pad(htm, ((0, 0), (radius, radius), (radius, radius)), mode='symmetric')
        # htm = torch.from_numpy(htm).float()
        # oks_kernel = torch.from_numpy(oks_kernel).float().to(htm.device).reshape(1, diameter, diameter)
        oks_kernel = oks_kernel.reshape(1, diameter, diameter)
        htm_conv = np.zeros_like(htm)
        for b in range(B):
            htm_conv[b, :, :] = convolve2d(htm[b, :, :], oks_kernel[b, :, :], mode='same', boundary='symm')
        # htm_conv = F.conv2d(htm.unsqueeze(1), oks_kernel.unsqueeze(1), padding='same')
        # htm_conv = htm_conv[:, :, radius:-radius, radius:-radius]
        htm_conv = htm_conv.reshape(-1, 1, H, W)
        heatmaps_covolved[:, k, :, :] = htm_conv

    
    heatmaps_covolved = heatmaps_covolved.reshape(B*K, H*W)
    y_locs, x_locs = np.unravel_index(
        np.argmax(heatmaps_covolved, axis=1), shape=(H, W))
    locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)

    # Apply mean-shift to get sub-pixel locations
    locs = _get_subpixel_maximums(heatmaps_covolved.reshape(B*K, H, W), locs)
    # breakpoint()


    # heatmaps_sums = heatmaps_flatten.sum(axis=(1, 2))
    # norm_heatmaps = heatmaps_flatten.copy()
    # norm_heatmaps[heatmaps_sums > 0] = heatmaps_flatten[heatmaps_sums > 0] / heatmaps_sums[heatmaps_sums > 0, None, None]
    

    # # Compute Parzen window with Gaussian blur along the edge instead of simple mirroring
    # x_pad = int(parzen_size * W + 0.5)
    # y_pad = int(parzen_size * H + 0.5)
    # # x_pad = 0
    # # y_pad = 0
    # kernel_size = int(min(H, W)*parzen_size + 0.5)
    # if kernel_size % 2 == 0:
    #     kernel_size += 1
    # # norm_heatmaps_pad_blur = np.pad(norm_heatmaps, ((0, 0), (x_pad, x_pad), (y_pad, y_pad)), mode='symmetric')
    # norm_heatmaps_pad = np.pad(norm_heatmaps, ((0, 0), (y_pad, y_pad), (x_pad, x_pad)), mode='constant', constant_values=0)
    # norm_heatmaps_pad_blur = gaussian_blur(norm_heatmaps_pad, kernel=kernel_size)
        
    # # norm_heatmaps_pad_blur[:, x_pad:-x_pad, y_pad:-y_pad] = norm_heatmaps
    
    # norm_heatmaps_pad_sum = norm_heatmaps_pad_blur.sum(axis=(1, 2))
    # norm_heatmaps_pad_blur[norm_heatmaps_pad_sum>0] = norm_heatmaps_pad_blur[norm_heatmaps_pad_sum>0] / norm_heatmaps_pad_sum[norm_heatmaps_pad_sum>0, None, None]
    
    # # # Save the blurred heatmaps
    # # for i in range(heatmaps.shape[0]):
    # #     tmp_htm = norm_heatmaps_pad_blur[i].copy()
    # #     tmp_htm = (tmp_htm - tmp_htm.min()) / (tmp_htm.max() - tmp_htm.min())
    # #     tmp_htm = (tmp_htm*255).astype(np.uint8)
    # #     tmp_htm = cv2.cvtColor(tmp_htm, cv2.COLOR_GRAY2BGR)
    # #     tmp_htm = cv2.applyColorMap(tmp_htm, cv2.COLORMAP_JET)

    # #     tmp_htm2 = norm_heatmaps_pad[i].copy()
    # #     tmp_htm2 = (tmp_htm2 - tmp_htm2.min()) / (tmp_htm2.max() - tmp_htm2.min())
    # #     tmp_htm2 = (tmp_htm2*255).astype(np.uint8)
    # #     tmp_htm2 = cv2.cvtColor(tmp_htm2, cv2.COLOR_GRAY2BGR)
    # #     tmp_htm2 = cv2.applyColorMap(tmp_htm2, cv2.COLORMAP_JET)

    # #     tmp_htm = cv2.addWeighted(tmp_htm, 0.5, tmp_htm2, 0.5, 0)

    # #     cv2.imwrite(f'heatmaps_blurred_{i}.png', tmp_htm)

    # # norm_heatmaps_pad = np.pad(norm_heatmaps, ((0, 0), (x_pad, x_pad), (y_pad, y_pad)), mode='edge')

    # y_idx, x_idx = np.indices(norm_heatmaps_pad_blur.shape[1:])

    # # breakpoint()
    # x_locs = np.sum(norm_heatmaps_pad_blur * x_idx, axis=(1, 2)) - x_pad
    # y_locs = np.sum(norm_heatmaps_pad_blur * y_idx, axis=(1, 2)) - y_pad
    
    # # mean_idx = np.argmax(heatmaps_flatten, axis=1)
    # # x_locs, y_locs = np.unravel_index(mean_idx, shape=(H, W))
    # # locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
    # # breakpoint()
    # # vals = heatmaps_flatten[np.arange(heatmaps_flatten.shape[0]), mean_idx]
    # # locs[vals <= 0.] = -1

    # # mean_idx = np.argmax(norm_heatmaps, axis=1)
    # # y_locs, x_locs = np.unravel_index(
    # #     mean_idx, shape=(H, W))
    
    # locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
    # # vals = np.amax(heatmaps_flatten, axis=1)
    
    
    x_locs_int = np.round(x_locs).astype(int)
    x_locs_int = np.clip(x_locs_int, 0, W-1)
    y_locs_int = np.round(y_locs).astype(int)
    y_locs_int = np.clip(y_locs_int, 0, H-1)
    vals = heatmaps_flatten[np.arange(B), np.arange(K), y_locs_int, x_locs_int]
    # breakpoint()
    # locs[vals <= 0.] = -1

    # print(mean_idx)
    # print(x_locs)
    # print(y_locs)
    # print(locs)
    heatmaps_covolved = heatmaps_covolved.reshape(B, K, H, W)

    if B > 1:
        locs = locs.reshape(B, K, 2)
        vals = vals.reshape(B, K)
        heatmaps_covolved = heatmaps_covolved.reshape(B, K, H, W)
    else:
        locs = locs.reshape(K, 2)
        vals = vals.reshape(K)
        heatmaps_covolved = heatmaps_covolved.reshape(K, H, W)

    if return_heatmap:
        return locs, vals, heatmaps_covolved
    else:
        return locs, vals       



def _get_subpixel_maximums(heatmaps, locs):
    # Extract integer peak locations
    x_locs = locs[:, 0].astype(np.int32)
    y_locs = locs[:, 1].astype(np.int32)

    # Ensure we are not near the boundaries (avoid boundary issues)
    valid_mask = (x_locs > 0) & (x_locs < heatmaps.shape[2] - 1) & \
                 (y_locs > 0) & (y_locs < heatmaps.shape[1] - 1)

    # Initialize the output array with the integer locations
    subpixel_locs = locs.copy()

    if np.any(valid_mask):
        # Extract valid locations
        x_locs_valid = x_locs[valid_mask]
        y_locs_valid = y_locs[valid_mask]

        # Compute gradients (dx, dy) and second derivatives (dxx, dyy)
        dx = (heatmaps[valid_mask, y_locs_valid, x_locs_valid + 1] - 
              heatmaps[valid_mask, y_locs_valid, x_locs_valid - 1]) / 2.0
        dy = (heatmaps[valid_mask, y_locs_valid + 1, x_locs_valid] - 
              heatmaps[valid_mask, y_locs_valid - 1, x_locs_valid]) / 2.0
        dxx = heatmaps[valid_mask, y_locs_valid, x_locs_valid + 1] + \
              heatmaps[valid_mask, y_locs_valid, x_locs_valid - 1] - \
              2 * heatmaps[valid_mask, y_locs_valid, x_locs_valid]
        dyy = heatmaps[valid_mask, y_locs_valid + 1, x_locs_valid] + \
              heatmaps[valid_mask, y_locs_valid - 1, x_locs_valid] - \
              2 * heatmaps[valid_mask, y_locs_valid, x_locs_valid]

        # Avoid division by zero by setting a minimum threshold for the second derivatives
        dxx = np.where(dxx != 0, dxx, 1e-6)
        dyy = np.where(dyy != 0, dyy, 1e-6)

        # Calculate the sub-pixel shift
        subpixel_x_shift = -dx / dxx
        subpixel_y_shift = -dy / dyy

        # Update subpixel locations for valid indices
        subpixel_locs[valid_mask, 0] += subpixel_x_shift
        subpixel_locs[valid_mask, 1] += subpixel_y_shift

    return subpixel_locs