jonigata/PoseTweak

1import numpy as np2import cv23 4def pil2cv(image):5    ''' PIL型 -> OpenCV型 '''6    new_image = np.array(image, dtype=np.uint8)7    if new_image.ndim == 2:  # モノクロ8        pass9    elif new_image.shape[2] == 3:  # カラー10        new_image = cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR)11    elif new_image.shape[2] == 4:  # 透過12        new_image = cv2.cvtColor(new_image, cv2.COLOR_RGBA2BGRA)13    return new_image14 15def candidate_to_json_string(arr):16    a = [f'[{x:.2f}, {y:.2f}]' for x, y, *_ in arr]17    return '[' + ', '.join(a) + ']'18 19# make subset to json20def subset_to_json_string(arr):21    arr_str = ','.join(['[' + ','.join([f'{num:.2f}' for num in row]) + ']' for row in arr])22    return '[' + arr_str + ']'23 24keypoint_index_mapping = [25    0, 26    17,27    6,28    8,29    10,30    5,31    7,32    9,33    12,34    14,35    16,36    11,37    13,38    15,39    2,40    1,41    4,42    3,43]44 45def convert_keypoints(keypoints):46    return [keypoints[i] for i in keypoint_index_mapping]47 48def convert_to_openpose(pose_result):49    candidate = []50    subset = []51    for d in pose_result:52        n = len(candidate)53        if d['bbox'][4] < 0.9: 54            continue55        keypoints = d['keypoints'][:, :2].tolist()56        midpoint = [(keypoints[5][0] + keypoints[6][0]) / 2, (keypoints[5][1] + keypoints[6][1]) / 2]57        keypoints.append(midpoint)58        candidate.extend(convert_keypoints(keypoints))59        m = len(candidate)60        subset.append([j for j in range(n, m)])61 62    return candidate, subset63 64