action_demo_with_video.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 smutils.utils_data import save_pkl_data
import shutil
 
#cfg 경로 설정
det_cfg_path = '/workspace/smlab/smdetect/configs/yolo/yolov8.py' # yolov8
track_cfg_path = '/workspace/smlab/smtrack/configs/bytetrack/bytetrack_base.py'
pose_cfg_path = '/workspace/smlab/smpose/configs/mmpose/hrnet_trt.py' # hrnet
action_cfg_path = '/workspace/InnoTest/models/posec3d_action.py'
posture_cfg_path = '/workspace/InnoTest/models/posec3d_pose.py' 
 
# 모델 빌드
device = 'cuda:0'
det_model = smrunner.build_model(det_cfg_path)
track_model = smrunner.build_model(track_cfg_path)
pose_model = smrunner.build_model(pose_cfg_path)
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()
 
#샘플링 파라미터
#TODO: cfg로 옮겨서 runner의 파라미터로 설정할 필요가 있음 또는 pipline에 설정
action_sample = 20
pose_sample = 6
 
#레이블 맵 로드
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')
 
#레이블맵 변환 테이블
cvt_lower_map = [0, 1, 2, 0, 4, 5, 0, 7, 8, 9, 0, 11]
cvt_upper_map = [0, 0, 0, 3, 0, 0]
cvt_pose_map = [0, 0, 0, 0, 4, 4, 6, 0]
cvt_hand_map = [0, 0, 0, 3]
 
#시각화 파라미터
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'  
)
 
#비디오 경로 설정
video_folder = '/media/safemotion/HDD5/pjm_test/ai_park_test_video'
_, video_path_list = search_file(video_folder, '.mp4')
 
#결과 저장 위치
save_folder = '/media/safemotion/HDD5/pjm_test/ai_park_test_fps20'
create_directory(save_folder)
 
#임시 폴더
tmp_save_folder = '/media/safemotion/HDD5/pjm_test/tmp_clip_images_fps20'
create_directory(tmp_save_folder)
shutil.rmtree(tmp_save_folder)
 
#데모
for video_idx, video_path in enumerate(video_path_list):
    
    #임시 폴더 생성
    create_directory(tmp_save_folder)
 
    #저장 경로 설정
    video_name = video_path.split('/')[-1]
    save_video_path = os.path.join(save_folder, video_name)
 
    #비디오 로드
    cap = cv2.VideoCapture(video_path)
 
    if not cap.isOpened():
        print(f"Error: Could not open video({video_path}).")
        continue
 
    #비디오 정보 출력
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    print(f'{video_path} : {total_frames} frames, {fps} FPS')
 
    #행동인식 범위 설정, 행동은 2초 구간 포즈는 0.5초 구간
    action_k = int(fps+0.5)*2
    pose_k = int(fps+0.5)//2
 
    #초기화
    track_data = {} #추적 객체 저장 변수, 추적아이디별 박스, 스켈레톤 정보 저장
    frame_id = 0 #프레임 아이디, 0부터 1씩 증가
    track_model.tracker.reset() #트래커 리셋, 여러 동영상을 하나의 트래커로 사용하기 때문에 동영상 시작전에 리셋함
    while True:
        #비디오에서 이미지 로드
        ret, image = cap.read()
 
        #마지막 프레임
        if not ret:
            print("Reached the end of the video or encountered an error.")
            break
 
        #진행 정도 출력
        print(f'{video_name} : {frame_id+1} / {total_frames}', end='\r')
        img_shape = image.shape[:2]
 
        #가시화 이미지 생성
        vis_img = image.copy()
 
        # 모델 inference(검출, 추적, 포즈)
        det_result = det_model.run_detector(image)
        track_result = track_model.run_tracker(det_result['det_bboxes'], det_result['det_labels'], frame_id)
        pose_result = pose_model.run_detector(image, track_result['track_bboxes'][0])
        
        #추적 아이디별로 포즈 데이터 저장
        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
            
            #포스처 모델 inference
            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]
 
        #프레임 아이디 증가
        frame_id+=1
        
        #시각화 이미지 저장
        name = f'{frame_id:09d}.jpg'
        save_path = os.path.join(tmp_save_folder, name)
        cv2.imwrite(save_path, vis_img)
 
    # 시각화 이미지 비디오로 저장
    make_video(tmp_save_folder, save_video_path, fps=fps, half=False) #동영상 생성
    shutil.rmtree(tmp_save_folder) #임시 폴더 삭제