pose_transform.py

Go to the documentation of this file.

 
# PoseDecode
#  -> 키포인트, 스코어 샘플링
 
# PoseCompact
#  -> 키포인트의 바운더리 찾고(마진 추가) -> 키포인트 좌표 변환(바운더리 왼쪽 상단을 0, 0)
#  -> img_shape를 (바운더리의 높이,너비)로 변경
#  -> 자르는 영역 계산해둠
 
# Resize
#  -> img_shape를 64x64 로 바꿈
#  -> sacale_factor 계산함
#  -> sacale_factor를 키포인트에 곱함
 
# RandomResizedCrop
#  -> 크롭 영역 설정
#  -> img_shape를 크롭된 크기로 설정
#  -> 키포인트 크롭 : 키포인트 좌표에서 크롭 영역의 시작점(좌상단)을 빼줌
 
# Resize
#  ->   56x56
 
# Flip
#  -> 키포인트 좌우 반전하고, 좌우 키포인트 및 스코어를 인덱스에 맞도록 바꿔줌
 
import numpy as np
import random
 
 
def pose_sampling(sample, clip_len):
    """
    스켈레톤 데이터를 랜덤으로 샘플링하는 기능
    args:
        sample (dict): 데이터, 여러 프레임의 스켈레톤 좌표와 스코어 값이 저장된 'keypoint', 'keypoint_score'를 포함하여야함
            'image_feature'가 존재하면 'image_feature'도 함께 샘플링함
        clip_len (int): 샘플링 수량
    return (dict): 샘플링된 데이터
    """
 
    num_frames = sample['keypoint'].shape[1] #프레임 수
    if num_frames < clip_len: #프레임 수가 clip_len보다 작을 경우
        start = np.random.randint(0, num_frames)
        inds = np.arange(start, start + clip_len)
    elif clip_len <= num_frames < 2 * clip_len: #프레임수가 clip_len의 2배수보다 작을경우
        basic = np.arange(clip_len)
        inds = np.random.choice(
            clip_len + 1, num_frames - clip_len, replace=False)
        offset = np.zeros(clip_len + 1, dtype=np.int32)
        offset[inds] = 1
        offset = np.cumsum(offset)
        inds = basic + offset[:-1]
    else: #프레임수가 clip_len의 2배수보다 클 경우
        bids = np.array(
            [i * num_frames // clip_len for i in range(clip_len + 1)])
        bsize = np.diff(bids)
        bst = bids[:clip_len]
        offset = np.random.randint(bsize)
        inds = bst + offset
    
    sample['keypoint'] = sample['keypoint'][:, inds].astype(np.float32)
    sample['keypoint_score'] = sample['keypoint_score'][:, inds].astype(np.float32)
    if 'image_feature' in sample:
        sample['image_feature'] = sample['image_feature'][:, inds]
    return sample
 
 
def pose_compact(sample):
    """
    모든 스켈레톤을 포함하는 영역을 계산하고 좌상단의 좌표가 (0, 0)이 되도록 스켈레톤 좌표를 shifting함
    스켈레톤 좌표의 전체 영역의 크기(높이, 너비)는 'kp_hw'에 저장됨
    args:
        sample (dict): 데이터, 스켈레톤 좌표가 저장된 'keypoint'를 포함함
    return (dict): shifting함된 스켈레톤 좌표와 영역의 크기가 포함된 데이터
    """
    keypoint = sample['keypoint']
    # keypoint shape : (1, 60, 17, 2) 
    kp_x = keypoint[..., 0] #x 좌표
    kp_y = keypoint[..., 1] #y 좌표
 
    #좌표의 최대 최소값
    min_x = kp_x.min()
    min_y = kp_y.min()
    max_x = kp_x.max()
    max_y = kp_y.max()
 
    #패딩(마진) 추가
    padding = 0.25
    center = ((max_x + min_x) / 2, (max_y + min_y) / 2) #박스 중심
    half_width = (max_x - min_x) / 2 * (1 + padding) #패딩이 추가된 너비 절반
    half_height = (max_y - min_y) / 2 * (1 + padding) #패딩이 추가된 높이 절반
    
    #너비 높이 중 최대값을 사용
    max_val = max(half_width, half_height) 
    half_width = max_val
    half_height = max_val
 
    #스켈레톤 영역(패딩 추가)
    min_x, max_x = center[0] - half_width, center[0] + half_width
    min_y, max_y = center[1] - half_height, center[1] + half_height
 
    #타입변환
    min_x, min_y = int(min_x), int(min_y)
    max_x, max_y = int(max_x), int(max_y)
 
    #스켈레톤 영역의 좌상단이 0, 0이 되도록 좌표 변환(shifting)
    kp_x -= min_x
    kp_y -= min_y
 
    #영역 크기 설정
    sample['kp_hw'] = (max_y - min_y, max_x - min_x)
    
    return sample
 
def pose_shift(sample, shift_ratio=0.05):
    """
    스켈레톤 좌표를 일정 비율만큼 이동 시키는 기능
    args:
        sample (dict): 데이터, 스켈레톤 좌표가 저장된 'keypoint'를 포함함
        shift_ratio (float): 좌표 이동 비율
    return (dict): 좌표가 이동된 스켈레톤을 포함하는 데이터
    """
 
    keypoint = sample['keypoint']
    #좌표의 최대, 최소 설정
    x1 = keypoint[:, 0, :, 0].min()
    x2 = keypoint[:, 0, :, 0].max()
    y1 = keypoint[:, 0, :, 1].min()
    y2 = keypoint[:, 0, :, 1].max()
 
    #이동 비율 설정
    w_range = (x2-x1)*shift_ratio
    h_range = (y2-y1)*shift_ratio
 
    #이동할 값을 랜덤으로 설정
    dx = random.uniform(-w_range, w_range)
    dy = random.uniform(-h_range, h_range)
 
    #좌표 변환
    sample['keypoint'][:, :, :, 0] += dx
    sample['keypoint'][:, :, :, 1] += dy
 
    return sample
 
def pose_resize(sample, scale=(64,64)):
    """
    스켈레톤 전체 영역의 크기를 지정한 크기로 변환
    args:
        sample (dict): 데이터, 스켈레톤 좌표가 저장된 'keypoint'와 스켈레톤 영역의 크기가 저장된 'kp_hw'를 포함함
        scale (tuple): 변환하려는 크기, (높이, 너비)의 형태
    return (dict): 리사이즈된 스켈레톤을 포함한 데이터
    """
 
    keypoint = sample['keypoint']
    h, w = sample['kp_hw'] #현재 스켈레톤 영역의 크기
    new_h, new_w = scale #리사이즈될 영역
 
    #리사이즈(좌표 변환)
    scale_factor = np.array([new_w / w, new_h / h], dtype=np.float32) #스케일 계산
    sample['keypoint'] = keypoint * scale_factor #변환
    sample['kp_hw'] = scale #영역 크기 설정
 
    return sample
 
 
def get_crop_bbox(kp_hw, area_range, aspect_ratio_range, max_attempts=10):
    """
    랜덤으로 자르는 영역을 생성하는 기능
    args:
        kp_hw (tuple): 스켈레톤을 모두 포함하는 박스의 영역 크기
        area_range (tuple): 자르는 영역 크기의 범위, 원본 대비 비율
        aspect_ratio_range (tuple): aspect_ratio(가로 세로 비율) 범위
        max_attempts (int): 랜덤 생성 최대 반복횟수, 10번안에 파라미터를 만족하는 영역이 생성되지 않으면 고정된 영역을 반환함
    return (float, float, float, float): left, top, right, bottom
    """
 
    #영역 크기 계산
    h, w = kp_hw
    area = h * w
 
    #타겟 aspect_ratio 및 영역 크기 설정
    min_ar, max_ar = aspect_ratio_range
    aspect_ratios = np.exp( np.random.uniform(np.log(min_ar), np.log(max_ar), size=max_attempts) )
    target_areas = np.random.uniform(*area_range, size=max_attempts) * area
 
    #aspect_ratio에 따른 자르는 너비 높이 후보 설정
    candidate_crop_w = np.round(np.sqrt(target_areas * aspect_ratios)).astype(np.int32)
    candidate_crop_h = np.round(np.sqrt(target_areas / aspect_ratios)).astype(np.int32)
 
    #조건에 맞는 영역이 있는지 검색, 조건에 맞는 영역이 있다면 해당 영역을 반환함
    for i in range(max_attempts):
        crop_w = candidate_crop_w[i]
        crop_h = candidate_crop_h[i]
 
        #자르는 영역의 크기가 기존 영역보다 작은지 체크
        if crop_h <= h and crop_w <= w:
            x_offset = random.randint(0, w - crop_w)
            y_offset = random.randint(0, h - crop_h)
            return x_offset, y_offset, x_offset + crop_w, y_offset + crop_h
 
    # Fallback
    crop_size = min(h, w)
    x_offset = (w - crop_size) // 2
    y_offset = (h - crop_size) // 2
    return x_offset, y_offset, x_offset + crop_size, y_offset + crop_size
 
def pose_random_crop(sample, area_range=(0.56, 1.0), aspect_ratio_range=(3 / 4, 4 / 3)):
    """
    키포인트를 랜덤으로 자르는 기능, 랜덤으로 영역을 생성하고 자름
    args:
        sample (dict): 데이터, 스켈레톤 좌표가 저장된 'keypoint'와 스켈레톤 영역의 크기가 저장된 'kp_hw'를 포함함
        area_range (tuple): 자르는 영역 크기의 범위, 원본 대비 비율
        aspect_ratio_range (tuple): aspect_ratio(가로 세로 비율) 범위
    return (dict): 잘린 키포인트가 포함된 데이터
    """
 
    keypoint = sample['keypoint']
    h, w = sample['kp_hw']
 
    #자르는 영역 설정
    left, top, right, bottom = get_crop_bbox( (h, w), area_range, aspect_ratio_range) 
    new_h, new_w = bottom - top, right - left
    crop_bbox = np.array([left, top, right, bottom])
 
    #자르기(좌표 변환)
    sample['keypoint'] = keypoint - crop_bbox[:2]
    sample['kp_hw'] = (new_h, new_w)
 
    return sample
 
def pose_flip(sample, flip_ratio=0.5):
    """
    스켈레톤을 좌우 반전하는 기능
    args:
        sample (dict): 데이터, 여러 프레임의 스켈레톤 좌표와 스코어 값이 저장된 'keypoint', 'keypoint_score'와 스켈레톤 영역의 크기가 저장된 'kp_hw'를 포함함
        flip_ratio (float): 좌우 반전 할 확률
    return (dict): 일정 확률로 좌우 반전이된 스켈레톤을 포함하는 데이터
    """
 
    #좌우 반전 진행 여부 설정
    flip = np.random.rand() < flip_ratio
 
    if flip: #좌우 반전 진행
        #17개 스켈레톤 좌표 인덱스, 코(0번) 제외
        left_kp = [1, 3, 5, 7, 9, 11, 13, 15] #왼쪽 스켈레톤
        right_kp = [2, 4, 6, 8, 10, 12, 14, 16] #오른쪽 스켈레톤
 
        h, w = sample['kp_hw']
        keypoint = sample['keypoint']
        keypoint_score = sample['keypoint_score']
 
        keypoint[..., 0] = w - keypoint[..., 0] #좌우 반전
        
        #스켈레톤 인덱스 변경
        new_order = list(range(keypoint.shape[2]))
        for left, right in zip(left_kp, right_kp):
            new_order[left] = right
            new_order[right] = left
        keypoint = keypoint[:, :, new_order]
        keypoint_score = keypoint_score[:, :, new_order]
 
        sample['keypoint'] = keypoint
        sample['keypoint_score'] = keypoint_score
 
    return sample
 
 
def generate_a_heatmap(pose_heatmap, centers, max_values):
    """
    하나의 스켈레톤에 대한 히트맵을 생성하는 기능, 여러명의 사람일 수 있음
    args:
        pose_heatmap (np.array): 1종 스켈레톤의 히트맵, shape (H, W)
        centers (np.array): 1-종의 스켈레톤 좌표(ex-손목, 코 등), shape (num_person, 2)
        max_values (np.array): 1-종의 스켈레톤 스코어, shape (num_person)
    return (): 
    """
    sigma = 0.6
    eps = 1e-4
    img_h, img_w = pose_heatmap.shape
    for center, max_value in zip(centers, max_values):
        
        if max_value < eps:
            continue
        
        #패치 영역 설정
        mu_x, mu_y = center[0], center[1]
        st_x = max(int(mu_x - 3 * sigma), 0)
        ed_x = min(int(mu_x + 3 * sigma) + 1, img_w)
        st_y = max(int(mu_y - 3 * sigma), 0)
        ed_y = min(int(mu_y + 3 * sigma) + 1, img_h)
        x = np.arange(st_x, ed_x, 1, np.float32)
        y = np.arange(st_y, ed_y, 1, np.float32)
 
        # if the keypoint not in the heatmap coordinate system
        if not (len(x) and len(y)):
            continue
        y = y[:, None]
 
        #패치의 값 셋팅, 가우시안 분포
        patch = np.exp(-((x - mu_x)**2 + (y - mu_y)**2) / 2 / sigma**2)
        patch = patch * max_value #스코어값 적용
 
        #여러 사람일 경우 중복된 위에 값이 있을 경우 최대값 사용
        pose_heatmap[st_y:ed_y, st_x:ed_x] = \
            np.maximum(pose_heatmap[st_y:ed_y, st_x:ed_x], patch)
 
def generate_heatmap(pose_heatmap, kps, kpscores):
    """
    한 프레임에 대한 스켈레톤 히트맵을 생성하는 기능
    args:
        pose_heatmap :
        kps : 스켈레톤 좌표, shape (num_person, 17, 2) 
        kpscores : 스켈레톤 스코어, shape (num_person, 17)
    """
    num_kp = kps.shape[1]
    for i in range(num_kp):
        generate_a_heatmap(pose_heatmap[i], kps[:, i], kpscores[:, i])
 
def make_pose_heatmap(sample):
    """
    스켈레톤 좌표를 사용해서 히트맵을 생성하는 기능
    args:
        sample (dict): 데이터, 여러 프레임의 스켈레톤 좌표와 스코어 값이 저장된 'keypoint', 'keypoint_score'와 스켈레톤 영역의 크기가 저장된 'kp_hw'를 포함함
    return (dict): 스켈레톤을 기반으로한 히트맵이 포함된 데이터, 히트맵의 키는 'pose_heatmap'임
    """
    keypoint = sample['keypoint']
    keypoint_score = sample['keypoint_score']
 
    h, w = sample['kp_hw']
    num_frame = keypoint.shape[1]
    num_c  = keypoint.shape[2]
    
    #히트맵 공간 미리 생성
    pose_heatmap = np.zeros([num_frame, num_c, h, w], dtype=np.float32)
    for i in range(num_frame):
        
        kps = keypoint[:, i]# shape (num_person, 17, 2), 17=스켈레톤 수
        kpscores = keypoint_score[:, i]
 
        #히트맵 생성
        generate_heatmap(pose_heatmap[i], kps, kpscores)
 
    # T, C, H, W -> C, T, H, W
    # T : 프레임, C : 스켈레톤 수
    pose_heatmap = np.transpose(pose_heatmap, (1, 0, 2, 3))
    sample['pose_heatmap'] = pose_heatmap
    return sample