update_anno_pose.py

Go to the documentation of this file.

import numpy as np
import mmcv
import copy
import cv2
import os
from tqdm import tqdm
 
from smdataset.AnnotationDataManager import AnnotationDataManager
from smutils.utils_os import search_file, create_directory
from smutils.utils_vis import vis_pose_coco_skeleton
import smrunner
 
def pose_inference_one_frame(model, image, bboxes):
    """
    한프레임에서 검출된 박스에대해 포즈 모델을 inference하는 기능
    args:
        model: smpose로 빌드한 포즈 모델
        image (np.array): 원본 이미지
        bboxes (np.array): 검출 박스(x,y,x,y), shape (num_person, 4 or 5 or 6)
    return:
        kp (np.array): 모든 사람에 대한 스켈레톤 좌표, shape (num_person, 17, 3)
        pose_result (list[dict]): 
            track_id (int): 추적 아이디
            bbox (np.array): 박스 위치(x, y, x, y)
            keypoints (np.array): 스켈레톤 좌표, shape (17, 3)
    """
    
    num_person = bboxes.shape[0]
    kp = np.zeros((num_person, 17, 3), dtype=np.float32)
 
    pose_result = model.run_detector(image, bboxes)
    for j, item in enumerate(pose_result):
        kp[j] = item['keypoints']
    
    return kp, pose_result
 
 
def update_annotation_pose_data(model, data_manager, save_pose_img_folder=None):
    """
    어노테이션 파일에 포즈모델로 예측한 스켈레톤을 업데이트하는 기능
    args:
        model: smpose로 빌드한 포즈 모델
        data_manager: AnnotationDataManager 객체, 어노테이션 데이터파일로 초기화까지 진행
        save_pose_img_folder: 스켈레톤 시각화 이미지를 저장할 폴더, 설정하지 않으면 시각화하지 않음
    """
 
    #시각화 이미지를 저장할 폴더 생성
    if save_pose_img_folder is not None:
        create_directory(save_pose_img_folder)
    
    #이미지 숫자 체크
    img_num = data_manager.get_number_of_image()
    for i in tqdm(range(img_num)):
        #이미지에 있는 어노테이션 데이터를 가져옴
        anno_data = data_manager.get_anno_data_in_image()
        #이미지 경로를 가져옴
        img_path = data_manager.get_image_path()
        
        #어노테이션 박스 구조 변경 (x, y, w, h) -> (x, y, x, y)
        bbox_in_image = []
        for anno in anno_data:
            bbox_tmp = copy.deepcopy(anno['bbox'])
            bbox_tmp[2] += bbox_tmp[0] #TODO: 영상 크기로 최대치 체크
            bbox_tmp[3] += bbox_tmp[1] #TODO: 영상 크기로 최대치 체크
            bbox_tmp[0] = max(bbox_tmp[0], 0)
            bbox_tmp[1] = max(bbox_tmp[1], 0)
            
            bbox_in_image.append(bbox_tmp)
 
        #박스가 없으면 다음 이미지로 넘어감
        if len(bbox_in_image) == 0:
            data_manager.move_image_right()
            continue
 
        bboxes = np.array(bbox_in_image) #박스 타입 변경
        image = cv2.imread(img_path) #이미지 로드
        key_points, pose_result = pose_inference_one_frame(model, image, bboxes) #포즈 inference
        
        #스켈레톤 데이터 업데이트 진행
        for anno, kp in zip(anno_data, key_points):
            kp = kp.tolist() #스켈레톤 데이터를 리스트로 변경(이중리스트가 됨)
            data_manager.update_annotation_data_by_anno_id(anno['id'], 'keypoints', kp) #업데이트
        
        #스켈레톤 시각화 이미지 저장
        if save_pose_img_folder is not None:
            #이미지에 스켈레톤 그리기
            image = vis_pose_coco_skeleton(image, pose_result)
            
            #저장 경로 생성
            file_name = img_path.split('/')[-1]
            save_path = os.path.join(save_pose_img_folder, file_name)
 
            #시각화 이미지 저장
            cv2.imwrite(save_path, image)
        
        #다음 이미지로 이동
        data_manager.move_image_right()
    
    #스켈레톤 업데이트된 어노테이션 데이터 저장
    data_manager.save_annotation()
    
 
def run_update_annotation_pose_data(args):
    """
    어노테이션 파일들에 포즈모델로 예측한 스켈레톤을 업데이트하는 기능
    args:
        save_json_root (str): 스켈레톤이 업데이트된 어노테이션을 저장할 폴더, 지정하지 않으면 원본 파일에 덮어씀
        pose_config (str): smpose 모델이 정의된 config 파일 경로
        dataset_folder_list (list[str]): 어노테이션 파일이 저장된 폴더 리스트
        skip (int): 스킵할 어노테이션 파일 수량
        save_pose_img_root (str): 스켈레톤 시각화 영상을 저장할 폴더, 설정하지 않으면 시각화 하지 않음
    """
    
    #json 저장 폴더 생성
    if args.save_json_root is not None:
        create_directory(args.save_json_root)
 
    #모델 빌드
    model = smrunner.build_model(args.pose_config)
    
    #어노테이션 관리자 생성
    data_manager = AnnotationDataManager()
    
    #어노테이션 리스트, 입력된 폴더에 존재하는 모든 어노테이션 파일을 리스트에 저장
    json_name_list = []
    json_path_list = []
    for dataset_folder in args.dataset_folder_list:
        name_list, path_list = search_file(dataset_folder, '.json') #dataset_folder에 저장된 어노테이션 파일 경로들을 찾음
        json_name_list.extend(name_list)
        json_path_list.extend(path_list)
    
    file_num = len(json_path_list) #파일 수량, 진행 정도 표시를 위함
    for i, (json_name, json_path) in enumerate(zip(json_name_list, json_path_list)):
        #스킵
        if i < args.skip:
            continue
        
        #저장 위치 설정
        save_path = None
        if args.save_json_root is not None:
            save_path = os.path.join(args.save_json_root, json_name)
        
        #어노테이션 관리자 초기화
        data_manager.load_annotation(json_path, save_path=save_path) #로드
        data_manager.init_annotation() #초기화
        
        #진행도 출력
        print(f"[ {i+1} / {file_num}] : load path : {json_path}")
        print(f"[ {i+1} / {file_num}] : save path : {data_manager.get_save_path()}")
 
        #저장 경로 설정
        save_pose_img_folder = None
        if args.save_pose_img_root is not None:
            folder = json_name.split('.')[0]
            save_pose_img_folder = os.path.join(args.save_pose_img_root, folder)
 
        #어노테이션 파일 하나에 대해서 스켈레톤 업데이트 진행
        update_annotation_pose_data(model, data_manager, save_pose_img_folder)
        print('')