SCAIL-Pose/pose_draw/reshape_utils.py

import numpy as np
import random

def pose_offset(offset_x, offset_y, pose, vitposes):
    for i in range(len(pose["bodies"]["candidate"])):
        bodies = pose["bodies"]
        faces = pose["faces"][i]
        right_hand = pose["hands"][2 * i]
        left_hand = pose["hands"][2 * i + 1]
        candidate = bodies["candidate"][i]
        subset = bodies["subset"][i]

        for face_point in faces:
            if face_point[0] == -1 and face_point[1] == -1:
                continue
            face_point[0] = face_point[0] + offset_x
            face_point[1] = face_point[1] + offset_y
        for left_hand_point in left_hand:
            if left_hand_point[0] == -1 and left_hand_point[1] == -1:
                continue
            left_hand_point[0] = left_hand_point[0] + offset_x
            left_hand_point[1] = left_hand_point[1] + offset_y
        for right_hand_point in right_hand:
            if right_hand_point[0] == -1 and right_hand_point[1] == -1:
                continue
            right_hand_point[0] = right_hand_point[0] + offset_x
            right_hand_point[1] = right_hand_point[1] + offset_y

        assert len(candidate) == len(subset), f"candidate, length: {len(candidate)} and subset, length: {len(subset)} must have the same length"
        for idx, body_point in enumerate(candidate):
            if subset[idx] == -1:
                continue
            if body_point[0] == -1 and body_point[1] == -1:
                continue
            body_point[0] = body_point[0] + offset_x
            body_point[1] = body_point[1] + offset_y

    if vitposes is not None:
        for vitpose in vitposes:
            vitpose['keypoints_body'] = vitpose['keypoints_body'] + np.array([offset_x, offset_y, 0])
            vitpose['keypoints_left_hand'] = vitpose['keypoints_left_hand'] + np.array([offset_x, offset_y, 0])
            vitpose['keypoints_right_hand'] = vitpose['keypoints_right_hand'] + np.array([offset_x, offset_y, 0])


def pose_whole_scale(scale_x, scale_y, pose, vitposes):   # 目前仅能用于单人
    # 获取最左端的x 最右端的x 最上端的y 最下端的y
    min_x = float('inf')
    max_x = float('-inf')
    min_y = float('inf')
    max_y = float('-inf')
    if len(pose["bodies"]["candidate"]) > 1:
        return
    for i in range(len(pose["bodies"]["candidate"])):
        bodies = pose["bodies"]
        faces = pose["faces"][i]
        right_hand = pose["hands"][2 * i]
        left_hand = pose["hands"][2 * i + 1]
        candidate = bodies["candidate"][i]
        subset = bodies["subset"][i]

        for face_point in faces:
            if face_point[0] == -1 and face_point[1] == -1:
                continue
            min_x = min(min_x, face_point[0])
            max_x = max(max_x, face_point[0])
            min_y = min(min_y, face_point[1])
            max_y = max(max_y, face_point[1])
        for left_hand_point in left_hand:
            if left_hand_point[0] == -1 and left_hand_point[1] == -1:
                continue
            min_x = min(min_x, left_hand_point[0])
            max_x = max(max_x, left_hand_point[0])
            min_y = min(min_y, left_hand_point[1])
            max_y = max(max_y, left_hand_point[1])
        for right_hand_point in right_hand:
            if right_hand_point[0] == -1 and right_hand_point[1] == -1:
                continue
            min_x = min(min_x, right_hand_point[0])
            max_x = max(max_x, right_hand_point[0])
            min_y = min(min_y, right_hand_point[1])
            max_y = max(max_y, right_hand_point[1])
        assert len(candidate) == len(subset), "candidate and subset must have the same length"
        for idx, body_point in enumerate(candidate):
            if subset[idx] == -1:
                continue
            if body_point[0] == -1 and body_point[1] == -1:
                continue
            min_x = min(min_x, body_point[0])
            max_x = max(max_x, body_point[0])
            min_y = min(min_y, body_point[1])
            max_y = max(max_y, body_point[1])

    # 计算缩放比例
    # 根据bbox中心进行dilate, 倍数为x_scale, y_scale
    bbox_center_x = (min_x + max_x) / 2
    bbox_center_y = (min_y + max_y) / 2
    for i in range(len(pose["bodies"]["candidate"])):
        bodies = pose["bodies"]
        candidate = bodies["candidate"][i]
        for body_point in candidate:
            if body_point[0] == -1 and body_point[1] == -1:
                continue
            body_point[0] = body_point[0] + (body_point[0] - bbox_center_x) * scale_x   # scale越大，越远离中心，变化越大
            body_point[1] = body_point[1] + (body_point[1] - bbox_center_y) * scale_y
        for face_point in faces:
            if face_point[0] == -1 and face_point[1] == -1:
                continue
            face_point[0] = face_point[0] + (face_point[0] - bbox_center_x) * scale_x
            face_point[1] = face_point[1] + (face_point[1] - bbox_center_y) * scale_y
        for left_hand_point in left_hand:
            if left_hand_point[0] == -1 and left_hand_point[1] == -1:
                continue
            left_hand_point[0] = left_hand_point[0] + (left_hand_point[0] - bbox_center_x) * scale_x
            left_hand_point[1] = left_hand_point[1] + (left_hand_point[1] - bbox_center_y) * scale_y
        for right_hand_point in right_hand:
            if right_hand_point[0] == -1 and right_hand_point[1] == -1:
                continue
            right_hand_point[0] = right_hand_point[0] + (right_hand_point[0] - bbox_center_x) * scale_x
            right_hand_point[1] = right_hand_point[1] + (right_hand_point[1] - bbox_center_y) * scale_y

    if vitposes is not None:
        min_x = float('inf')
        max_x = float('-inf')
        min_y = float('inf')
        max_y = float('-inf')
        if len(vitposes) > 1:
            return
        for vitpose in vitposes:
            for body_point in vitpose['keypoints_body']:
                if body_point[2] < 0.4:
                    continue
                min_x = min(min_x, body_point[0])
                max_x = max(max_x, body_point[0])
                min_y = min(min_y, body_point[1])
                max_y = max(max_y, body_point[1])
            for left_hand_point in vitpose['keypoints_left_hand']:
                if left_hand_point[2] < 0.4:
                    continue
                min_x = min(min_x, left_hand_point[0])
                max_x = max(max_x, left_hand_point[0])
                min_y = min(min_y, left_hand_point[1])
                max_y = max(max_y, left_hand_point[1])
            for right_hand_point in vitpose['keypoints_right_hand']:
                if right_hand_point[2] < 0.4:
                    continue
                min_x = min(min_x, right_hand_point[0])
                max_x = max(max_x, right_hand_point[0])
                min_y = min(min_y, right_hand_point[1])
                max_y = max(max_y, right_hand_point[1])
            bbox_center_x = (min_x + max_x) / 2
            bbox_center_y = (min_y + max_y) / 2
            for body_point in vitpose['keypoints_body']:
                body_point[0] = body_point[0] + (body_point[0] - bbox_center_x) * scale_x
                body_point[1] = body_point[1] + (body_point[1] - bbox_center_y) * scale_y
            for left_hand_point in vitpose['keypoints_left_hand']:
                left_hand_point[0] = left_hand_point[0] + (left_hand_point[0] - bbox_center_x) * scale_x
                left_hand_point[1] = left_hand_point[1] + (left_hand_point[1] - bbox_center_y) * scale_y
            for right_hand_point in vitpose['keypoints_right_hand']:
                right_hand_point[0] = right_hand_point[0] + (right_hand_point[0] - bbox_center_x) * scale_x
                right_hand_point[1] = right_hand_point[1] + (right_hand_point[1] - bbox_center_y) * scale_y




# 如果做多人增强，需要修改逻辑，每个人都有随机性，现在都是做单人增强，增强方法一致
def pose_reshape(alpha, pose, vitposes,
                 anchor_point, anchor_part, end_point, end_part, 
                 affected_body_points, affected_faces_points, affected_left_hands_points, affected_right_hands_points):
    """
    直接修改原始pose中的某些点的位置，通过根据一个不动点和端点之间的向量偏移被影响点的位置。
    一次只能有一个anchor_point, 一个end_point，多个affected_points
    alpha: 变化比例
    anchor_point: 起始点
    anchor_part: 起始点所属的身体部分，0代表body candidate, 1代表face, 2代表hand
    end_point: 中止点
    end_part: 中止点所属的身体部分，0代表body candidate, 1代表face, 2代表hand
    """
    for i in range(len(pose["bodies"]["candidate"])):
        bodies = pose["bodies"]
        # faces = pose["faces"][i:i+1]
        faces = pose["faces"][i]
        # hands = pose["hands"][2*i:2*i+2]
        right_hand = pose["hands"][2 * i]
        left_hand = pose["hands"][2 * i + 1]
        candidate = bodies["candidate"][i]
        # subset = bodies["subset"][i:i+1]
        subset = bodies["subset"][i]
        assert anchor_part == 0, "anchor part must belong to the body"
        anchor_x, anchor_y = candidate[anchor_point]
        if subset[anchor_point] == -1:
            continue

        if end_part == 0:
            end_x, end_y = candidate[end_point]
            if subset[end_point] == -1:
                continue
        elif end_part == 1:
            end_x, end_y = faces[end_point]
        # 不考虑Hands
        if end_x == -1 and end_y == -1:
            continue
        
        vector_x = end_x - anchor_x
        vector_y = end_y - anchor_y
        offset_x = vector_x * alpha
        offset_y = vector_y * alpha
        
        for affected_body_point in affected_body_points:
            if subset[affected_body_point] == -1:
                continue
            affected_x, affected_y = candidate[affected_body_point]
            candidate[affected_body_point] = [affected_x + offset_x, affected_y + offset_y]
        for affected_faces_point in affected_faces_points:
            affected_x, affected_y = faces[affected_faces_point]
            if affected_x == -1 and affected_y == -1:
                continue
            faces[affected_faces_point] = [affected_x + offset_x, affected_y + offset_y]
        for affected_hands_point in affected_left_hands_points:
            affected_x, affected_y = left_hand[affected_hands_point]
            if affected_x == -1 and affected_y == -1:
                continue
            left_hand[affected_hands_point] = [affected_x + offset_x, affected_y + offset_y]
        for affected_hands_point in affected_right_hands_points:
            affected_x, affected_y = right_hand[affected_hands_point]
            if affected_x == -1 and affected_y == -1:
                continue
            right_hand[affected_hands_point] = [affected_x + offset_x, affected_y + offset_y]

    if vitposes is not None:
        vit_visible_threshold = 0.5
        for vitpose in vitposes:
            assert anchor_part == 0, "anchor part must belong to the body"
            anchor_x, anchor_y, _ = vitpose['keypoints_body'][anchor_point]
            if vitpose['keypoints_body'][anchor_point][2] < vit_visible_threshold:
                continue
            if end_part == 0:
                end_x, end_y, _ = vitpose['keypoints_body'][end_point]
                if vitpose['keypoints_body'][end_point][2] < vit_visible_threshold:
                    continue
            # 不考虑hands
            elif end_part == 1:
                continue       # ViTPose的脸不动

            vector_x = end_x - anchor_x
            vector_y = end_y - anchor_y
            offset_x = vector_x * alpha
            offset_y = vector_y * alpha

            for affected_body_point in affected_body_points:
                if vitpose['keypoints_body'][affected_body_point][2] < vit_visible_threshold:
                    continue
                vitpose['keypoints_body'][affected_body_point][0] = vitpose['keypoints_body'][affected_body_point][0] + offset_x
                vitpose['keypoints_body'][affected_body_point][1] = vitpose['keypoints_body'][affected_body_point][1] + offset_y
            # 左右标记相反，需要互换
            for affected_left_hand_point in affected_left_hands_points:
                if vitpose['keypoints_right_hand'][affected_left_hand_point][2] < vit_visible_threshold:
                    continue
                vitpose['keypoints_right_hand'][affected_left_hand_point][0] = vitpose['keypoints_right_hand'][affected_left_hand_point][0] + offset_x
                vitpose['keypoints_right_hand'][affected_left_hand_point][1] = vitpose['keypoints_right_hand'][affected_left_hand_point][1] + offset_y
            for affected_right_hand_point in affected_right_hands_points:
                if vitpose['keypoints_left_hand'][affected_right_hand_point][2] < vit_visible_threshold:
                    continue
                vitpose['keypoints_left_hand'][affected_right_hand_point][0] = vitpose['keypoints_left_hand'][affected_right_hand_point][0] + offset_x
                vitpose['keypoints_left_hand'][affected_right_hand_point][1] = vitpose['keypoints_left_hand'][affected_right_hand_point][1] + offset_y
            


# reshapePool只负责形变，骨骼偏移、丢弃等得从draw层来做
class reshapePool:
    def __init__(self, alpha):  # 对每个视频只初始化一次
        self.faces_indices = np.arange(0, 68)
        self.left_hands_indices = np.arange(0, 21)
        self.right_hands_indices = np.arange(0, 21)
        self.alpha = alpha  # 0.1
        self.offset_x = random.uniform(-1/16, 1/16)
        self.offset_y = random.uniform(-1/16, 1/16)
        self.scale_x = random.uniform(-alpha/2, alpha/2)
        self.scale_y = random.uniform(-alpha/2, alpha/2)

        self.body_reshape_methods = [
            self.extend_body,
            self.extend_arm,
            self.extend_leg,
            self.shrink_body,
            self.shrink_arm,
            self.shrink_leg,
        ]
        self.scale_reshape_methods = [
            self.offset_wholebody,
            # self.scale_wholebody,
            self.dilate_face,
            self.shrink_face,
        ]
        self.selected_methods = random.sample(self.body_reshape_methods, 2) + random.sample(self.scale_reshape_methods, random.choice([0, 1]))

    def apply_random_reshapes(self, pose, vitposes=None):
        # Apply the two selected reshape methods
        for method in self.selected_methods:
            method(pose, vitposes)

    def offset_wholebody(self, pose, vitposes):
        pose_offset(self.offset_x, self.offset_y, pose, vitposes)

    def scale_wholebody(self, pose, vitposes):
        pose_whole_scale(self.scale_x, self.scale_y, pose, vitposes)


    def extend_body(self, pose, vitposes):
        pose_reshape(self.alpha, pose, vitposes,
                     1, 0, 8, 0,
                     [8, 9, 10], [], [], []
                     )
        pose_reshape(self.alpha, pose, vitposes,
                    1, 0, 11, 0,
                    [11, 12, 13], [], [], []
                    )
        
    def shrink_body(self, pose, vitposes):
        pose_reshape(-self.alpha, pose, vitposes,
                     1, 0, 8, 0,
                     [8, 9, 10], [], [], []
                     )
        pose_reshape(-self.alpha, pose, vitposes,
                    1, 0, 11, 0,
                    [11, 12, 13], [], [], []
                    )
        
    def extend_arm(self, pose, vitposes):
        pose_reshape(self.alpha, pose, vitposes,
                     2, 0, 3, 0,
                     [3, 4], [], self.left_hands_indices, []
                     )
        pose_reshape(self.alpha, pose, vitposes,
                    5, 0, 6, 0,
                    [6, 7], [], [], self.right_hands_indices
                    )
        pose_reshape(self.alpha, pose, vitposes,
                     3, 0, 4, 0,
                     [4], [], self.left_hands_indices, []
                     )
        pose_reshape(self.alpha, pose, vitposes,
                    6, 0, 7, 0,
                    [7], [], [], self.right_hands_indices
                    )
            
    def shrink_arm(self, pose, vitposes):
        pose_reshape(-self.alpha, pose, vitposes,
                     2, 0, 3, 0,
                     [3, 4], [], self.left_hands_indices, []
                     )
        pose_reshape(-self.alpha, pose, vitposes,
                    5, 0, 6, 0,
                    [6, 7], [], [], self.right_hands_indices
                    )
        pose_reshape(-self.alpha, pose, vitposes,
                     3, 0, 4, 0,
                     [4], [], self.left_hands_indices, []
                     )
        pose_reshape(-self.alpha, pose, vitposes,
                    6, 0, 7, 0,
                    [7], [], [], self.right_hands_indices
                    )

    def extend_leg(self, pose, vitposes):
        pose_reshape(self.alpha, pose, vitposes,
                     8, 0, 9, 0,
                     [9, 10], [], [], []
                     )
        pose_reshape(self.alpha, pose, vitposes,
                    11, 0, 12, 0,
                    [12, 13], [], [], []
                    )
        pose_reshape(self.alpha, pose, vitposes,
                     9, 0, 10, 0,
                     [10], [], [], []
                     )
        pose_reshape(self.alpha, pose, vitposes,
                    12, 0, 13, 0,
                    [13], [], [], []
                    )
        
    def shrink_leg(self, pose, vitposes):
        pose_reshape(-self.alpha, pose, vitposes,
                     8, 0, 9, 0,
                     [9, 10], [], [], []
                     )
        pose_reshape(-self.alpha, pose, vitposes,
                    11, 0, 12, 0,
                    [12, 13], [], [], []
                    )
        pose_reshape(-self.alpha, pose, vitposes,
                     9, 0, 10, 0,
                     [10], [], [], []
                     )
        pose_reshape(-self.alpha, pose, vitposes,
                    12, 0, 13, 0,
                    [13], [], [], []
                    )

    def extend_shoulder(self, pose, vitposes):
        pose_reshape(self.alpha, pose, vitposes,
                     1, 0, 2, 0,
                     [2, 3, 4], [], self.left_hands_indices, []
                     )
        pose_reshape(self.alpha, pose, vitposes,
                    1, 0, 5, 0,
                    [5, 6, 7], [], [], self.right_hands_indices
                    )
        
    def shrink_shoulder(self, pose, vitposes):
        pose_reshape(-self.alpha, pose, vitposes,
                     1, 0, 2, 0,
                     [2, 3, 4], [], self.left_hands_indices, []
                     )
        pose_reshape(-self.alpha, pose, vitposes,
                    1, 0, 5, 0,
                    [5, 6, 7], [], [], self.right_hands_indices
                    )
        
    
    def dilate_face(self, pose, vitposes):
        for i in self.faces_indices:
            pose_reshape(self.alpha, pose, vitposes,
                     0, 0, i, 1,
                     [], [i], [], []
                     )
        for i in [14, 15, 16, 17]:
            pose_reshape(self.alpha, pose, vitposes,
                    0, 0, i, 0,
                    [i], [], [], []
                    )
        
    def shrink_face(self, pose, vitposes):    # 缩小脸会看不清
        for i in self.faces_indices:
            pose_reshape(-self.alpha / 2, pose, vitposes,
                     0, 0, i, 1,
                     [], [i], [], []
                     )
            
        for i in [14, 15, 16, 17]:
            pose_reshape(-self.alpha / 2, pose, vitposes,
                    0, 0, i, 0,
                    [i], [], [], []
                    )