generate_clip.py

Go to the documentation of this file.

 
import os
from smutils.utils_os import search_file, create_directory, save_json
from smdataset.AnnotationDataManager import AnnotationDataManager
from smdataset.gen_data.cvt_labelmap import *
import copy
 
def search_clip_list_v21(anno_list, image_ids, frame_num):
    """
    어노테이션 데이터(한명에 해당하는 데이터)에서 행동별로 클립을 생성하는 기능
    v21 어노테이션 버전에 대해서만 동작함.
    args:
        anno_list (list[dict]): 어노테이션 데이터 리스트(한명에 대한 정보), action_id에 action 키가 존재 해야함, 해당 키에 해당 프레임에서의 클래스 정보가 입력되어 있어야함   
        image_ids (list(int)): 어노테이션 데이터의 이미지 아이디 리스트
        frame_num (int): 어노테이션 데이터의 총 프레임 수
    return (dict): 카테고리별 클래스의 클립 리스트, 클립은 클립의 시작 프레임 인덱스(start), 끝 프레임 인덱스(end), 카테고리의 클래스 번호(label), 중복방지를 위한 순번(seq)을 포함하고 있음
    """
    #액션 클래스 저장
    action_list = []
    for anno in anno_list:
        action_label = anno['action_id']['action'] #행동 클래스
 
        #행동 클래스가 리스트에 없으면 행동 클래스 저장
        if action_label not in action_list:
            action_list.append(action_label)
    
    #클래스별 클립 생성
    clip_list = []
    for label in action_list:
        #최대 간격 설정
        max_interval = 15 if label == 0 else 30
        
        #클립의 이미지 인덱스 저장 변수
        idx_list = [[]]
        
        #클립 정보 생성, 클래스에 해당하는 이미지 인덱스 리스트
        for anno in anno_list:
            action_label = anno['action_id']['action']
            image_idx = image_ids.index(anno['image_id'])
            
            #체크하는 레이블과 다르면 패스
            if label != action_label:
                continue
            
            #클립이 비어있으면 생성
            if len(idx_list[-1]) == 0:
                idx_list[-1].append(image_idx)
            
            #프레임 사이 간격이 max_interval 이하면 이어서 저장, 아니면 새로운 클립 생성
            if image_idx - idx_list[-1][-1] < max_interval:
                idx_list[-1].append(image_idx)
            else:
                idx_list.append([image_idx])
        
        #클립 생성
        for i, idxs in enumerate(idx_list):
 
            #클립의 시작, 끝 초기화
            start, end = idxs[0], idxs[-1]
 
            # #행동 구간이 7프레임 이하이면 패스
            # if end - start < 7:
            #     continue
            
            #마진 추가, 클래스가 0(행동 없음)이 아니면 마진 추가함
            if label != 0:
                #클립을 최소 90프레임으로 설정함
                clip_len = end - start #클립 길이
                margin = max( (90 - clip_len)//2, 30) #마진값 설정
                start = max(0, start - margin) #마진이 포함된 시작 프레임 설정
                end = min(frame_num-1, max(start + 60, end + margin)) #마진이 포함된 끝 프레임 설정
            
            #학습용 클립 길이기반 유효성 검사
            if end - start < 60:
                continue
            
            #클립 생성
            clip = dict(start=start, end=end, label=label, seq=i)
            clip_list.append(clip)
            
    return clip_list
 
def search_label_list(anno_list, cat_list):
    """
    어노테이션 데이터에 존재하는 카테고리별 클래스 리스트 탐색기능
    args:
        anno_list (list[dict]): 어노테이션 데이터 리스트, 데이터에는 카테고리별 클래스가 설정된 'action_id'키를 포함하여야함
        cat_list (list[str]): 카테고리 리스트
    return (dict): 어노테이션 데이터에 존재하는 각 카테고리별 클래스 리스트
    """
 
    #카테고리별 클래스 리스트 초기화
    ret_dict = dict()
    for cat in cat_list:
        ret_dict[cat] = []
    
    #탐색
    for anno in anno_list:
        for cat in cat_list:
            label = anno['action_id'][cat]
            if label not in ret_dict[cat]:
                ret_dict[cat].append(label)
 
    return ret_dict
 
 
def search_clip_list_v22(anno_list, image_ids, frame_num, action_train_frame=60, pose_train_frame=10):
    """
    어노테이션 데이터(한명에 해당하는 데이터)에서 행동별로 클립을 생성하는 기능
    args:
        anno_list (list[dict]): 어노테이션 데이터 리스트(한명에 대한 정보), 데이터에는 카테고리별 클래스가 설정된 'action_id'키를 포함하여야함
        image_ids (list(int)): 어노테이션 데이터의 이미지 아이디 리스트
        frame_num (int): 어노테이션 데이터의 총 프레임 수
        action_train_frame (int): 행동 인식의 학습 데이터 생성에 사용되는 프레임 수
        pose_train_frame (int): 포즈 인식의 학습 데이터 생성에 사용되는 프레임 수
    return (dict): 카테고리별 클래스의 클립 리스트, 클립은 클립의 시작 프레임 인덱스(start), 끝 프레임 인덱스(end), 카테고리의 클래스 번호(label), 중복방지를 위한 순번(seq)을 포함하고 있음
    """
 
    #클립 생성 파라미터 설정
    cat_info = dict()
    cat_list=['action_upper', 'action_lower', 'pose', 'hand', 'foot']
    for cat in cat_list:
        cat_info[cat] = dict()
 
    cat_info['action_upper']['info'] = cvt_labelmap_v22_action_upper
    cat_info['action_lower']['info'] = cvt_labelmap_v22_action_lower
    cat_info['pose']['info'] = cvt_labelmap_v22_action_pose
    cat_info['hand']['info'] = cvt_labelmap_v22_action_hand
    cat_info['foot']['info'] = cvt_labelmap_v22_action_foot
    cat_info['action_upper']['train_frame'] = action_train_frame
    cat_info['action_lower']['train_frame'] = action_train_frame
    cat_info['pose']['train_frame'] = pose_train_frame
    cat_info['hand']['train_frame'] = pose_train_frame
    cat_info['foot']['train_frame'] = pose_train_frame
 
    #어노테이션 데이터에서 카테고리별 클래스 리스트를 탐색, 데이터에 존재하는 모든 클래스를 의미함
    # label_info 예시
    # {'action_upper': [0],
    # 'action_lower': [0],
    # 'pose': [1],
    # 'hand': [0, 1],
    # 'foot': [0]}
    label_info = search_label_list(anno_list, cat_list)
 
    #카테고리별 클래스에 해당하는 이미지 인덱스 리스트, 클립을 구성하는 프레임 인덱스
    # clip_info init 예시
    # {'action_upper': {0: [[]]},
    # 'action_lower': {0: [[]]},
    # 'pose': {1: [[]]},
    # 'hand': {0: [[]], 1: [[]]},
    # 'foot': {0: [[]]}}
    clip_info = dict()
    for cat in cat_list:
        clip_info[cat] = dict()
        label_list = label_info[cat]
        for label in label_list:
            clip_info[cat][label] = [[]]
    
    #카테고리별 클래스에 해당하는 이미지 인덱스 리스트
    clip_interval = 15 #프레임 간격의 임계치, clip_interval이상이면 다른 클립으로 분리함
    for anno in anno_list:
        image_idx = image_ids.index(anno['image_id']) #어노테이션 데이터의 이미지 인덱스
 
        #카테고리별 클립 정보 설정
        for cat in cat_list:
            label = anno['action_id'][cat]
            
            if len(clip_info[cat][label][-1]) == 0: #클립이 비어있으면 초기화
                clip_info[cat][label][-1].append(image_idx)
            
            if image_idx - clip_info[cat][label][-1][-1] < clip_interval: #프레임 간격이 임계치보다 작으면 같은 클립으로 취급함
                clip_info[cat][label][-1].append(image_idx)
            else:
                clip_info[cat][label].append([image_idx]) #프레임 간격이 임계치보다 크면 새로운 클립으로 생성함
 
    #카테고리별 클래스의 클립 설정
    # clip_list init 예시
    # {'action_upper': {0: []},
    # 'action_lower': {0: []},
    # 'pose': {1: []},
    # 'hand': {0: [], 1: []},
    # 'foot': {0: []}}
    clip_list = dict()
    for cat in cat_list:
        clip_list[cat] = dict()
        label_list = label_info[cat]
        for label in label_list:
            clip_list[cat][label] = []
 
    #카테고리별 클래스의 클립 생성
    for cat, cat_clip_info in clip_info.items():
        for label, clip_frame_list in cat_clip_info.items():
 
            #클래스의 유효성 체크
            if label == -1 or label is None:
                continue
 
            #클립 정보 생성
            seq = 0
            for frame_list in clip_frame_list:
 
                #클립의 시작, 끝 초기화
                start, end = frame_list[0], frame_list[-1]
 
                #클립 길이 유효성 체크(마진 없는 클립 길이)
                clip_len = end - start +1
                if clip_len < cat_info[cat]['info'][label][3]:
                    continue
                
                #마진 추가
                if cat_info[cat]['info'][label][4]: #사용자 설정 마진
                    #클립을 최소 90프레임으로 설정함
                    margin = max( (90 - clip_len)//2, 30) #마진값 설정
                    start = max(0, start - margin) #마진이 포함된 시작 프레임 설정
                    end = min(frame_num-1, max(start + 60, end + margin)) #마진이 포함된 끝 프레임 설정
 
                #학습용 클립 길이기반 유효성 검사
                if (end - start +1) < cat_info[cat]['train_frame']:
                    continue
                
                #클립 생성
                clip = dict(start=start, end=end, label=label, seq=seq)
                clip_list[cat][label].append(clip)
                seq += 1
 
    return clip_list
 
 
def save_clip_v21(clip_list, input_anno_list, image_ids, base_anno, save_folder, json_name):
    """
    한사람에 대해 행동별로 생성된 클립을 저장하는 기능
    v21 어노테이션 버전에 대해서만 동작함 
    args:
        clip_list (list[dict]): 카테고리별 클래스의 클립정보(한명에 대한 정보), 클립의 시작, 끝, 클래스 저장되어 있음
        input_anno_list (list[dict]): 한사람에 대한 어노테이션 데이터
        image_ids (list[str]): 이미지 아이디 리스트
        base_anno : 어노테이션 데이터 원본, 파일을 읽었을때의 초기 상태
        save_folder (str): 클립 데이터를 저장할 폴더
        json_name (str): 원본 클립의 이름
    """
 
    #이미지 인덱스와 입력 어노테이션 데이터의 인덱스를 매칭시켜줌 
    idx_link = {}
    for i, anno in enumerate(input_anno_list): 
        image_idx = image_ids.index(anno['image_id'])
        idx_link[image_idx] = i
 
    #클립 데이터 저장
    for clip in clip_list:
        start = clip['start']
        end = clip['end']
        label = clip['label']
        seq =  clip['seq']
 
        #하나의 클립에 대한 정보만 저장
        anno_list = []
        for image_idx in range(start, end+1):
            #이미지 인덱스에 해당하는 어노테이션 데이터를 가져오기, 어노테이션 데이터가 없을 경우 None으로 셋팅
            anno = input_anno_list[idx_link[image_idx]] if image_idx in idx_link else None
 
            #어노테이션 데이터가 존재할 경우만 클립 정보 저장
            if anno is not None:
                anno_list.append(anno)
        
        #클립 데이터 생성, 원본어노테이션 파일에 클립 시작, 끝, 데이터 셋팅
        base_anno['annotations'] = anno_list
        base_anno['info']['start'] = start
        base_anno['info']['end'] = end
 
        #클립 데이터 저장
        track_id = anno_list[0]['track_id'] #추적 아이디
        new_name = json_name.replace('.json', f'T{track_id:02d}L{label:02d}S{seq:02d}.json') #클립 이름 설정
        save_path = os.path.join(save_folder, f'{label:02d}', new_name) #저장 경로
        save_json(base_anno, save_path) #저장
        print('  saved : ', save_path)
 
def save_clip_v22(clip_list, input_anno_list, image_ids, base_anno, save_folder, json_name):
    """
    한사람에 대해 행동별로 생성된 클립을 저장하는 기능
    args:
        clip_list (dict): 카테고리별 클래스의 클립정보(한명에 대한 정보)
            구조
                {
                    '카테고리 명': {클래스 번호: [ {클립정보(start, end, lable, seq)} ]}
                }
            구조 예시 : 
                {'action_upper': {0: [{'start': 20, 'end': 60, 'label':0, 'seq':0}]},
                'action_lower': {0: []},
                'pose': {1: []},
                'hand': {0: [], 1: []},
                'foot': {0: []}}
        input_anno_list (list[dict]): 한사람에 대한 어노테이션 데이터
        image_ids (list[str]): 이미지 아이디 리스트
        base_anno : 어노테이션 데이터 원본, 파일을 읽었을때의 초기 상태
        save_folder (str): 클립 데이터를 저장할 폴더
        json_name (str): 원본 클립의 이름
    """
 
    #이미지 인덱스와 입력 어노테이션 데이터의 인덱스를 매칭시켜줌 
    idx_link = {}
    for i, anno in enumerate(input_anno_list): 
        image_idx = image_ids.index(anno['image_id'])
        idx_link[image_idx] = i
 
    #클립 데이터 저장
    for cat, label_clip_info in clip_list.items(): #cat(str):카테고리, label_clip_info(dict):클래스에 대한 클립 리스트
        for label, clip_list in label_clip_info.items(): #label(int):클래스, clip_list(list[dict]):클립 리스트
            for clip in clip_list:
                start = clip['start']
                end = clip['end']
                label = clip['label']
                seq =  clip['seq']
 
                #하나의 클립에 대한 정보만 저장
                anno_list = []
                for image_idx in range(start, end+1):
                    #이미지 인덱스에 해당하는 어노테이션 데이터를 가져오기, 어노테이션 데이터가 없을 경우 None으로 셋팅
                    anno = input_anno_list[idx_link[image_idx]] if image_idx in idx_link else None
 
                    #어노테이션 데이터가 존재할 경우만 클립 정보 저장
                    if anno is not None:
                        anno_list.append(anno)
 
                #클립 데이터 생성, 원본어노테이션 파일에 클립 시작, 끝, 데이터 셋팅
                base_anno['annotations'] = anno_list
                base_anno['info']['start'] = start
                base_anno['info']['end'] = end
 
                #클립 데이터 저장
                track_id = anno_list[0]['track_id'] #추적 아이디
                new_name = json_name.replace('.json', f'T{track_id:02d}L{label:02d}S{seq:02d}.json') #클립 이름 설정
                save_path = os.path.join(save_folder, cat, f'{label:02d}', new_name) #저장 경로
                save_json(base_anno, save_path) #저장
                print('  saved : ', save_path)
 
def generate_clip(json_folder, save_folder):
    """
    어노테이션 데이터들을 개별 행동 클립으로 분할하는 기능
    args:
        json_folder (str): 어노테이션 파일이 저장된 폴더
        save_folder (str): 클립 데이터를 저장할 폴더
    """
 
    #어노테이션 파일 리스트
    json_name_list, json_path_list = search_file(json_folder, '.json')
 
    #데이터 관리자 생성
    data_manager = AnnotationDataManager()
 
    #클립 생성 프로세스
    N= len(json_name_list)
    for i, (json_name, json_path) in enumerate(zip(json_name_list, json_path_list)):
        print(f'{i}/{N} {json_path}')
 
        #데이터 관리자 초기화
        data_manager.load_annotation(json_path)
        data_manager.init_annotation()
        
        #원본 파일 복사, info와 images 키에 대한 정보가 필요함, 데이터를 클립의 데이터로 교체해서 저장하기 위함
        base_anno = copy.deepcopy(data_manager.dataset)
 
        #변수 설정
        labelmap_version = int(base_anno['info']['label_map_version']*10) #어노테이션 버전
        frames = data_manager.get_number_of_image() #프레임 수
        
        #추적 아이디 기준으로 구조 변경
        track_data = data_manager.get_track_anno_data(filter=True)
        
        #클립 생성
        for track_id, anno_list in track_data.items():
            if labelmap_version == 21: #버전 2.1
                simple_anno_list = cvt_labelmap_v21_pose_and_action(anno_list, cvt_labelmap_v21_action_to_simple_ver2, cvt_labelmap_v21_pose_to_simple_ver2) #학습을 위한 행동 라벨로 변경   
                clip_list = search_clip_list_v21(simple_anno_list, data_manager.image_ids, frames) #변경된 라벨 기준으로 클립 정보 추출
                save_clip_v21(clip_list, simple_anno_list, data_manager.image_ids, base_anno, save_folder, json_name) #클립 저장, 저장은 원본 행동 라벨 데이터로 저장함
            elif labelmap_version == 22: #버전 2.2
                simple_anno_list = cvt_labelmap_v22(anno_list) #학습을 위한 행동 라벨로 변경
                clip_list = search_clip_list_v22(simple_anno_list, data_manager.image_ids, frames) #변경된 라벨 기준으로 클립 정보 추출
                save_clip_v22(clip_list, anno_list, data_manager.image_ids, base_anno, save_folder, json_name) #클립 저장, 저장은 원본 행동 라벨 데이터로 저장함