controlnet_aux/open_pose/body.py

import math
from typing import List, NamedTuple, Union

import cv2
import numpy as np
import torch
from scipy.ndimage.filters import gaussian_filter

from . import util
from .model import bodypose_model


class Keypoint(NamedTuple):
 x: float
 y: float
 score: float = 1.0
 id: int = -1


class BodyResult(NamedTuple):
 # Note: Using `Union` instead of `|` operator as the ladder is a Python
 # 3.10 feature.
 # Annotator code should be Python 3.8 Compatible, as controlnet repo uses
 # Python 3.8 environment.
 # https://github.com/lllyasviel/ControlNet/blob/d3284fcd0972c510635a4f5abe2eeb71dc0de524/environment.yaml#L6
 keypoints: List[Union[Keypoint, None]]
 total_score: float
 total_parts: int


class Body(object):
 def __init__(self, model_path):
 self.model = bodypose_model()
 model_dict = util.transfer(self.model, torch.load(model_path))
 self.model.load_state_dict(model_dict)
 self.model.eval()

 def to(self, device):
 self.model.to(device)
 return self

 def __call__(self, oriImg):
 device = next(iter(self.model.parameters())).device
 # scale_search = [0.5, 1.0, 1.5, 2.0]
 scale_search = [0.5]
 boxsize = 368
 stride = 8
 padValue = 128
 thre1 = 0.1
 thre2 = 0.05
 multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
 heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
 paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))

 for m in range(len(multiplier)):
 scale = multiplier[m]
 imageToTest = util.smart_resize_k(oriImg, fx=scale, fy=scale)
 imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
 im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
 im = np.ascontiguousarray(im)

 data = torch.from_numpy(im).float()
 data = data.to(device)
 # data = data.permute([2, 0, 1]).unsqueeze(0).float()
 with torch.no_grad():
 Mconv7_stage6_L1, Mconv7_stage6_L2 = self.model(data)
 Mconv7_stage6_L1 = Mconv7_stage6_L1.cpu().numpy()
 Mconv7_stage6_L2 = Mconv7_stage6_L2.cpu().numpy()

 # extract outputs, resize, and remove padding
 # heatmap = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[1]].data), (1, 2, 0)) # output 1 is heatmaps
 heatmap = np.transpose(np.squeeze(Mconv7_stage6_L2), (1, 2, 0)) # output 1 is heatmaps
 heatmap = util.smart_resize_k(heatmap, fx=stride, fy=stride)
 heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
 heatmap = util.smart_resize(heatmap, (oriImg.shape[0], oriImg.shape[1]))

 # paf = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[0]].data), (1, 2, 0)) # output 0 is PAFs
 paf = np.transpose(np.squeeze(Mconv7_stage6_L1), (1, 2, 0)) # output 0 is PAFs
 paf = util.smart_resize_k(paf, fx=stride, fy=stride)
 paf = paf[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
 paf = util.smart_resize(paf, (oriImg.shape[0], oriImg.shape[1]))

 heatmap_avg += heatmap_avg + heatmap / len(multiplier)
 paf_avg += + paf / len(multiplier)

 all_peaks = []
 peak_counter = 0

 for part in range(18):
 map_ori = heatmap_avg[:, :, part]
 one_heatmap = gaussian_filter(map_ori, sigma=3)

 map_left = np.zeros(one_heatmap.shape)
 map_left[1:, :] = one_heatmap[:-1, :]
 map_right = np.zeros(one_heatmap.shape)
 map_right[:-1, :] = one_heatmap[1:, :]
 map_up = np.zeros(one_heatmap.shape)
 map_up[:, 1:] = one_heatmap[:, :-1]
 map_down = np.zeros(one_heatmap.shape)
 map_down[:, :-1] = one_heatmap[:, 1:]

 peaks_binary = np.logical_and.reduce(
 (one_heatmap >= map_left, one_heatmap >= map_right, one_heatmap >= map_up, one_heatmap >= map_down, one_heatmap > thre1))
 peaks = list(zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0])) # note reverse
 peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks]
 peak_id = range(peak_counter, peak_counter + len(peaks))
 peaks_with_score_and_id = [peaks_with_score[i] + (peak_id[i],) for i in range(len(peak_id))]

 all_peaks.append(peaks_with_score_and_id)
 peak_counter += len(peaks)

 # find connection in the specified sequence, center 29 is in the position 15
 limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
 [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
 [1, 16], [16, 18], [3, 17], [6, 18]]
 # the middle joints heatmap correpondence
 mapIdx = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44], [19, 20], [21, 22], \
 [23, 24], [25, 26], [27, 28], [29, 30], [47, 48], [49, 50], [53, 54], [51, 52], \
 [55, 56], [37, 38], [45, 46]]

 connection_all = []
 special_k = []
 mid_num = 10

 for k in range(len(mapIdx)):
 score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
 candA = all_peaks[limbSeq[k][0] - 1]
 candB = all_peaks[limbSeq[k][1] - 1]
 nA = len(candA)
 nB = len(candB)
 indexA, indexB = limbSeq[k]
 if (nA != 0 and nB != 0):
 connection_candidate = []
 for i in range(nA):
 for j in range(nB):
 vec = np.subtract(candB[j][:2], candA[i][:2])
 norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
 norm = max(0.001, norm)
 vec = np.divide(vec, norm)

 startend = list(zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
 np.linspace(candA[i][1], candB[j][1], num=mid_num)))

 vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
 for I in range(len(startend))])
 vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
 for I in range(len(startend))])

 score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
 score_with_dist_prior = sum(score_midpts) / len(score_midpts) + min(
 0.5 * oriImg.shape[0] / norm - 1, 0)
 criterion1 = len(np.nonzero(score_midpts > thre2)[0]) > 0.8 * len(score_midpts)
 criterion2 = score_with_dist_prior > 0
 if criterion1 and criterion2:
 connection_candidate.append(
 [i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2]])

 connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
 connection = np.zeros((0, 5))
 for c in range(len(connection_candidate)):
 i, j, s = connection_candidate[c][0:3]
 if (i not in connection[:, 3] and j not in connection[:, 4]):
 connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
 if (len(connection) >= min(nA, nB)):
 break

 connection_all.append(connection)
 else:
 special_k.append(k)
 connection_all.append([])

 # last number in each row is the total parts number of that person
 # the second last number in each row is the score of the overall configuration
 subset = -1 * np.ones((0, 20))
 candidate = np.array([item for sublist in all_peaks for item in sublist])

 for k in range(len(mapIdx)):
 if k not in special_k:
 partAs = connection_all[k][:, 0]
 partBs = connection_all[k][:, 1]
 indexA, indexB = np.array(limbSeq[k]) - 1

 for i in range(len(connection_all[k])): # = 1:size(temp,1)
 found = 0
 subset_idx = [-1, -1]
 for j in range(len(subset)): # 1:size(subset,1):
 if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
 subset_idx[found] = j
 found += 1

 if found == 1:
 j = subset_idx[0]
 if subset[j][indexB] != partBs[i]:
 subset[j][indexB] = partBs[i]
 subset[j][-1] += 1
 subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
 elif found == 2: # if found 2 and disjoint, merge them
 j1, j2 = subset_idx
 membership = ((subset[j1] >= 0).astype(int) + (subset[j2] >= 0).astype(int))[:-2]
 if len(np.nonzero(membership == 2)[0]) == 0: # merge
 subset[j1][:-2] += (subset[j2][:-2] + 1)
 subset[j1][-2:] += subset[j2][-2:]
 subset[j1][-2] += connection_all[k][i][2]
 subset = np.delete(subset, j2, 0)
 else: # as like found == 1
 subset[j1][indexB] = partBs[i]
 subset[j1][-1] += 1
 subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]

 # if find no partA in the subset, create a new subset
 elif not found and k < 17:
 row = -1 * np.ones(20)
 row[indexA] = partAs[i]
 row[indexB] = partBs[i]
 row[-1] = 2
 row[-2] = sum(candidate[connection_all[k][i, :2].astype(int), 2]) + connection_all[k][i][2]
 subset = np.vstack([subset, row])
 # delete some rows of subset which has few parts occur
 deleteIdx = []
 for i in range(len(subset)):
 if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
 deleteIdx.append(i)
 subset = np.delete(subset, deleteIdx, axis=0)

 # subset: n*20 array, 0-17 is the index in candidate, 18 is the total score, 19 is the total parts
 # candidate: x, y, score, id
 return candidate, subset
 
 @staticmethod
 def format_body_result(candidate: np.ndarray, subset: np.ndarray) -> List[BodyResult]:
 """
 Format the body results from the candidate and subset arrays into a list of BodyResult objects.
 
 Args:
 candidate (np.ndarray): An array of candidates containing the x, y coordinates, score, and id
 for each body part.
 subset (np.ndarray): An array of subsets containing indices to the candidate array for each
 person detected. The last two columns of each row hold the total score and total parts
 of the person.

 Returns:
 List[BodyResult]: A list of BodyResult objects, where each object represents a person with
 detected keypoints, total score, and total parts.
 """
 return [
 BodyResult(
 keypoints=[
 Keypoint(
 x=candidate[candidate_index][0],
 y=candidate[candidate_index][1],
 score=candidate[candidate_index][2],
 id=candidate[candidate_index][3]
 ) if candidate_index != -1 else None
 for candidate_index in person[:18].astype(int)
 ],
 total_score=person[18],
 total_parts=person[19]
 )
 for person in subset
 ]