action_demo_with_anno.py

Go to the documentation of this file.

 
#system path 설정
import sys
sys_path = ['/workspace/smlab', '/workspace']
for path in sys_path:
    if path not in sys.path:
        sys.path.append(path)
 
#import package
import os
import cv2
import smrunner
from smutils.utils_vis import draw_single_bbox_and_label
from smutils.utils_vis import vis_pose_coco_skeleton
from smutils.utils_os import search_file, create_directory
from smutils.utils_data import load_labelmap
from smutils.utils_video import make_video
 
from smdataset.AnnotationDataManager import AnnotationDataManager
 
import copy
import numpy as np
import shutil
 
#annotation data 변환 함수
def make_pose_result(anno_data):
    pose_result = []
    for anno in anno_data:
        tmp_data = {}
 
        box = copy.deepcopy(anno['bbox'])
 
        tmp_data['bbox'] = np.array(box)
        tmp_data['keypoints'] = np.array(anno['keypoints'])
        tmp_data['track_id'] = anno['track_id']
        pose_result.append(tmp_data)
 
    return pose_result
 
#시각화 파라미터
action_vis_param = dict(
    box_color = (0, 255, 0),
    box_thk = 3,
    txt_color = (255, 255, 255),
    txt_thk=3,
    txt_scale=1.5,
    box_type='xyxy'  
)
 
vis_param = dict(
    box_color = (0, 255, 0),
    box_thk = 3,
    txt_color = (255, 255, 255),
    txt_thk=3,
    txt_scale=1.5,
    box_type='xyxy'  
)
 
#경로 설정
action_cfg_path = '/workspace/InnoTest/models/posec3d_action.py'
posture_cfg_path = '/workspace/InnoTest/models/posec3d_pose.py'
image_base = '/media/safemotion/HDD2/SM_Dataset/action_2022/action_cctv_img/images' #이미지가 저장된 폴더
save_root = '/media/safemotion/HDD5/pjm_test/inno2023_result_select_half' #결과 동영상 저장 폴더
 
#어노테이션 파일이 저장된 폴더 리스트
anno_folder_list = [
    '/media/safemotion/HDD2/SM_Dataset/action_2022/action_cctv_img/annotations_action_v22_kid_1',
    '/media/safemotion/HDD2/SM_Dataset/action_2022/action_cctv_img/annotations_action_v22_kid_kkomo_1',
    '/media/safemotion/HDD2/SM_Dataset/action_2022/action_cctv_img/annotations_action_v22_kid_kkumjalam_1',
    '/media/safemotion/HDD2/SM_Dataset/action_2022/action_cctv_img/annotations_action_v22_kid_kkumjalam_hard_1',
    '/media/safemotion/HDD2/SM_Dataset/action_2022/action_cctv_img/annotations_action_v22_kid_kkomo_2',
    '/media/safemotion/HDD2/SM_Dataset/action_2022/action_cctv_img/annotations_action_v22_kid_kkumjalam_hard_1'
    '/media/safemotion/HDD2/SM_Dataset/action_2022/action_cctv_img/annotations_action_v22_2nd_3',
    '/media/safemotion/HDD2/SM_Dataset/action_2022/action_cctv_img/annotations_action_v22_2nd_4',              
]
 
#레이블맵 변환 테이블
cvt_lower_map = [0, 1, 2, 0, 4, 5, 6, 7, 8, 9, 0, 11]
cvt_upper_map = [0, 0, 0, 3, 4, 0]
cvt_pose_map = [0, 0, 0, 0, 4, 4, 6, 7]
cvt_hand_map = [0, 0, 0, 3]
 
#저장 폴더 생성
create_directory(save_root)
 
# 모델 빌드
device = 'cuda:0'
action_model = smrunner.build_model(action_cfg_path).to(device)
posture_model = smrunner.build_model(posture_cfg_path).to(device)
action_model.eval()
posture_model.eval()
 
#레이블 맵 로드
labelmap = dict()
labelmap['action_upper'] = load_labelmap('/workspace/smlab/smaction/datasets/safemotion_v22_upper_action.txt')
labelmap['action_lower'] = load_labelmap('/workspace/smlab/smaction/datasets/safemotion_v22_lower_action.txt')
labelmap['pose'] = load_labelmap('/workspace/smlab/smaction/datasets/safemotion_v22_pose.txt')
labelmap['hand'] = load_labelmap('/workspace/smlab/smaction/datasets/safemotion_v22_hand.txt')
 
#어노테이션 파일 경로 셋팅
json_path_list = []
for anno_folder in anno_folder_list:
    name_list, path_list = search_file(anno_folder, '.json')
    json_path_list.extend(path_list)
NUM = len(json_path_list)
 
#샘플링 파라미터 설정
#TODO: cfg로 옮겨서 runner의 파라미터로 설정할 필요가 있음 또는 pipline에 설정
video_fps = 30
action_k = video_fps*2
action_sample = 20
pose_k = video_fps//2
pose_sample = 6
 
#데이터 로더
data_manager = AnnotationDataManager()
 
#데모
for iter_n, json_path in enumerate(json_path_list):
    
    #저장경로 설정
    json_name = json_path.split('/')[-1]
    save_folder = os.path.join(save_root, json_name.split('.')[0]) #임시 저장 경로
    save_video_path = os.path.join(save_root, json_name.replace('.json', '.mp4')) #동영상 저장 경로
    create_directory(save_folder) #폴더 생성
 
    #데이터 매니저 초기화
    data_manager.load_annotation(json_path, image_base=image_base)
    data_manager.init_annotation()
 
    #초기화
    track_data = {} #추적 객체 저장 변수, 추적아이디별 박스, 스켈레톤 정보 저장
    frame_id = 0 #프레임 아이디, 0부터 1씩 증가
    while True:
        #진행도 출력
        print(f'[{iter_n+1}/{NUM}] {json_name} : {frame_id}', end='\r')
 
        #해당 프레임에 존재하는 어노테이션 데이터 읽기
        anno_data = data_manager.get_anno_data_in_image(filter=True)
 
        #이미지 로드
        image = data_manager.read_image()
 
        #가시화 이미지 생성
        vis_img = image.copy()
        vis_img = cv2.cvtColor(vis_img, cv2.COLOR_RGB2BGR)
 
        # 모델 inference(검출, 추적, 포즈) 결과와 동일한 포맷으로 만들기
        pose_result = make_pose_result(anno_data)
        
        #추적 아이디별로 포즈 데이터 저장
        for pose in pose_result: 
            track_id = int(pose['track_id'])
            if track_id not in track_data:
                track_data[track_id] = []
            pose['frame_id'] = frame_id
            track_data[track_id].append(pose)
        
        #스켈레톤 시각화    
        vis_img = vis_pose_coco_skeleton(vis_img, pose_result)
 
        delete_list = []
        txt_pos_y = 50
        for track_id, pose_q in track_data.items():
            #박스 시각화
            bbox = pose_q[-1]['bbox']
            
            last_frame = pose_q[-1]['frame_id']
            if last_frame == frame_id:
                label = f'{track_id:3d}'
                vis_img = draw_single_bbox_and_label(vis_img, bbox, label, **vis_param)
 
 
            #최근 1초동안 검출 안되면 제거
            if frame_id - last_frame > action_k*0.5:
                delete_list.append(track_id)
                continue
 
            #최근 15프레임동안 검출 안되면 넘김
            #포스처 때문에 넣은 조건, 하나의 모델로 변경되면 조건 변경 필요함
            if frame_id - last_frame > 15:
                continue
 
            #액션 모델 inference
            action_result = action_model.run_recognizer(pose_q, action_k, action_sample, device=device)
 
            #결과가 None면 유효성 체크에서 탈락한 것임
            if action_result is None:
                continue
            
            posture_result = posture_model.run_recognizer(pose_q[-pose_k:], pose_k, pose_sample, False, device=device)
            # #첫 프레임과 마지막 프레임 간격이 설정한 구간을 넘기면 맨 앞쪽 포즈 제거
            # #상황에 따라 루프로 제거해야 할 수도 있음
            if last_frame - pose_q[0]['frame_id'] >= action_k:
                del pose_q[0]
            
            
            #추론된 결과 변환
            pose_label = cvt_pose_map[posture_result['pred_pose']]
            upper_label = cvt_upper_map[action_result['pred_action_upper']]
            lower_label = cvt_lower_map[action_result['pred_action_lower']]
            hand_label = cvt_hand_map[posture_result['pred_hand']]
 
            #행동 라벨 및 박스 시각화
            label = f"{track_id:3d}: {labelmap['pose'][pose_label]}"
            
            if hand_label != 0:
                label += f"/{labelmap['hand'][hand_label]}"
            if upper_label != 0:
                label += f"/{labelmap['action_upper'][upper_label]}"
            if lower_label != 0:
                label += f"/{labelmap['action_lower'][lower_label]}"
            
            cv2.putText(vis_img, label, (10, txt_pos_y), cv2.FONT_HERSHEY_DUPLEX, 1.5, (0, 0, 0), 10, 1)
            cv2.putText(vis_img, label, (10, txt_pos_y), cv2.FONT_HERSHEY_DUPLEX, 1.5, (0, 255, 255), 2, 1)
            txt_pos_y += 45   
 
            vis_img = draw_single_bbox_and_label(vis_img, bbox, f'{track_id:3d}', **action_vis_param)
        
        #긴시간 미검출된 추적 아이디 제거
        for track_id in delete_list:
            del track_data[track_id]
 
        #마지막 프레임일경우 종료
        if data_manager.get_number_of_image() <= (data_manager.image_idx+1):
            break
 
        #다음프레임으로 이동
        data_manager.move_image_right()
 
        #프레임 아이디 증가
        frame_id+=1
 
        #시각화 이미지 저장
        name = f'{frame_id:06d}.jpg'
        save_path = os.path.join(save_folder, name)
        cv2.imwrite(save_path, vis_img)
 
    #시각화 이미지 비디오로 저장
    make_video(save_folder, save_video_path, fps=video_fps, half=True) #동영상 생성
    shutil.rmtree(save_folder) #임시 폴더 삭제