generate_clip_mtml.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_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, train_frame=60):
    """
    어노테이션 데이터(한명에 해당하는 데이터)에서 행동별로 클립을 생성하는 기능
    args:
        anno_list (list[dict]): 어노테이션 데이터 리스트(한명에 대한 정보), 데이터에는 카테고리별 클래스가 설정된 'action_id'키를 포함하여야함
        image_ids (list(int)): 어노테이션 데이터의 이미지 아이디 리스트
        frame_num (int): 어노테이션 데이터의 총 프레임 수
        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
 
    #어노테이션 데이터에서 카테고리별 클래스 리스트를 탐색, 데이터에 존재하는 모든 클래스를 의미함
    # 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 < 0 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 + train_frame, end + margin)) #마진이 포함된 끝 프레임 설정
 
                    #마진이 포함된 클립 길이 유효성 검사, 클립 길이 조건을 만족하지 못하면 시작 프레임을 조정함
                    if (end - start +1) < train_frame:
                        start = max(0, end - train_frame)
 
                else: #기본 마진 설정
                    margin = 15
                    start = max(0, start - margin)
                    end = min(frame_num-1, end + margin)
 
                #학습용 클립 길이기반 유효성 검사
                if (end - start +1) < 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_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)
 
        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():
            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) #클립 저장, 저장은 원본 행동 라벨 데이터로 저장함