util.py

import numpy as np
import cv2

def pil2cv(image):
    ''' PIL型 -> OpenCV型 '''
    new_image = np.array(image, dtype=np.uint8)
    if new_image.ndim == 2:  # モノクロ
        pass
    elif new_image.shape[2] == 3:  # カラー
        new_image = cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR)
    elif new_image.shape[2] == 4:  # 透過
        new_image = cv2.cvtColor(new_image, cv2.COLOR_RGBA2BGRA)
    return new_image

def candidate_to_json_string(arr):
    a = [f'[{x:.2f}, {y:.2f}]' for x, y, *_ in arr]
    return '[' + ', '.join(a) + ']'

# make subset to json
def subset_to_json_string(arr):
    arr_str = ','.join(['[' + ','.join([f'{num:.2f}' for num in row]) + ']' for row in arr])
    return '[' + arr_str + ']'

keypoint_index_mapping = [
    0, 
    17,
    6,
    8,
    10,
    5,
    7,
    9,
    12,
    14,
    16,
    11,
    13,
    15,
    2,
    1,
    4,
    3,
]

def convert_keypoints(keypoints):
    return [keypoints[i] for i in keypoint_index_mapping]

def convert_to_openpose(pose_result):
    candidate = []
    subset = []
    for d in pose_result:
        n = len(candidate)
        if d['bbox'][4] < 0.9: 
            continue
        keypoints = d['keypoints'][:, :2].tolist()
        midpoint = [(keypoints[5][0] + keypoints[6][0]) / 2, (keypoints[5][1] + keypoints[6][1]) / 2]
        keypoints.append(midpoint)
        candidate.extend(convert_keypoints(keypoints))
        m = len(candidate)
        subset.append([j for j in range(n, m)])

    return candidate, subset