generate_pkl_data_mtml.py

Go to the documentation of this file.

import pickle
import numpy as np
import copy
import os
 
from smdataset.gen_data.cvt_labelmap import *
from smdataset.AnnotationDataManager import AnnotationDataManager
 
from smutils.utils_os import search_file, create_directory
from smutils.utils_data import save_pkl_data, load_pkl_data, remove_items
import torch
import torch.nn as nn
 
def make_dummy_list(anno, num):
    """
    어노테이션 데이터를 복제해서 더미데이터 리스트를 생성하는 기능
    args:
        anno (dict): 복재하기 위한 어노테이션 데이터
        num (int): 복제 수량
    return (list[dict]): 더미데이터 리스트
    """
 
    #복사
    tmp_anno = copy.deepcopy(anno)
 
    #복사한 데이터의 키포인트 스코어를 0으로 셋팅함
    for i in range(len(tmp_anno['keypoints'])):
        tmp_anno['keypoints'][i][-1] = 0
    
    #유효성 셋팅
    tmp_anno['valid'] = 0
 
    dummy_list = []
    for i in range(num):
        dummy_list.append(copy.deepcopy(tmp_anno))
    return dummy_list
 
def check_invaild(anno_list, ratio):
    """
    데이터의 유효성 확인 기능
    args:
        anno_list (list[dict]): 확인하려는 데이터
            valid : 유효성 파라미터, 1 or 0, 1이면 유효한 데이터, 0이면 유효하지 않는 데이터
        ratio (float): 유효하지 않는 데이터 비율, 해당 비율보다 작으면 유효성 확인을 통과하지 못함
    return (bool): 유효성 확인 결과, 유효성 확인을 통과하지 못하면 True, 통과하면 False
    """
    sum_valid = 0
    for anno in anno_list:
        sum_valid += anno['valid']
 
    return sum_valid < len(anno_list)*ratio
 
 
def merge_pkl(folder, file_path_list=None ):
    """
    학습데이터 파일들을 하나로 합치는 기능
    args:
        folder (str): 학습데이터가 저장된 폴더
        file_path_list (list[str]) : 학습데이터 경로 리스트, 설정하지 않으면 folder에서 찾음
    return (list[dict]): 합쳐진 학습 데이터
    """
 
    if file_path_list == None:
        file_name_list, file_path_list = search_file(folder, fileEx='.pkl')
 
    merged_pkl_datas = []
    for path in file_path_list:
        with open(path, 'rb') as f:
            pkl_data = pickle.load(f)
            merged_pkl_datas += pkl_data
            print(f"{path} : {len(pkl_data)}")
            
    return merged_pkl_datas
 
def run_merge_pkl(args):
    """
    학습데이터 파일들을 하나로 합치고 저장하는 기능
    args:
        pkl_root (str): 학습데이터가 저장된 폴더
    """
    merged_pkl_datas = merge_pkl(args.pkl_root)
    print(f"total {args.merge_pkl_path} : {len(merged_pkl_datas)}")
    save_pkl_data(merged_pkl_datas, args.merge_pkl_path)
 
 
def check_label_v22(anno_list, category, target_label, ratio=0.3):
    """
    데이터에 특정 카테고리의 특정 클래스가 존재하는지 확인하는 기능
    args:
        anno_list (list[dict]): 데이터
            필수 데이터 : action_id에 라벨 이름(category)에 대한 키가 존재 해야함, 해당 키에 해당 프레임에서의 클래스 정보가 입력되어 있어야함         
        category (str): 데이터에서 확인하려는 카테고리
        target_label (int): 데이터에서 확인하려는 클래스
        ratio (float): 라벨 확인을 위한 데이터의 시작 위치(%), 클래스를 확인하는데 사용하는 전체 범위는 데이터 중앙의 (1-2*ratio) *100 %가 됨
        
    return (bool): 원하는 클래스 존재 유무
    """
 
    #확인 구간 설정
    l = len(anno_list)
    start_idx = int(l*ratio)
    end_idx = l-start_idx
    label_list = []
    
    #확인 구간의 클래스 라벨 리스트 생성, 빈도수를 확인하고 
    for i in range(start_idx, end_idx):
        anno = anno_list[i]
        
        #유효성 체크, 0 이하로 클래스가 저장되어 있으면 해당 프레임은 사용하지 않는 다는 의미임
        if anno['action_id'][category] < 0:
            continue
 
        #유효성 체크, valid가 0이면 더미데이터를 의미함
        if anno['valid'] == 0:
            continue
        
        label_list.append(anno['action_id'][category])
 
    return target_label in label_list
 
    # values, counts = np.unique(np.array(label_list), return_counts=True)
    # if target_label not in list(values):
    #     return False
    # return True
 
def make_label_v22(anno_list, category, category_margin_ratio):
    """
    데이터의 클래스를 결정하는 기능
    args:
        anno_list (list[dict]): 어노테이션 데이터 리스트, 각 항목은 'action_id'키를 포함해야 하며, 이키는 해당 프레임의 클래스 정보를 담고 있는 category(리스트) 키들을 포함해야 함
        category (list): 생성하려는 학습 데이터의 카테고리 리스트
        category_margin_ratio (dict): 카테고리별 라벨 결정을 위한 데이터의 시작 위치(%), 클래스를 결정하는데 사용하는 전체 범위는 데이터 중앙의 (1-2*ratio) *100 %가 됨
    return (dict): 카테고리별 데이터의 클래스, 유효성 검사나 클래스 확인이 어려운 경우 None을 반환함
    """
 
    l = len(anno_list) #데이터 전체 길이
 
    #카테고리별 클래스를 저장할 변수 초기화
    label_dict = dict()
    for cat in category:
        label_dict[cat] = []
 
    #설정구간에 대한 클래스 리스트를 생성함, 빈도수 확인을 위함
    for cat in category:
        ratio = category_margin_ratio[cat] #
        start_idx = int(l*ratio) #클래스를 결정할 구간의 시작 위치
        end_idx = l-start_idx #클래스를 결정할 구간의 끝 위치
        
        for i in range(start_idx, end_idx):
            anno = anno_list[i]
 
            #유효성 확인, 0보다 작으면 해당 정보는 사용하지 않는다는 의미임
            if anno['action_id'][cat] < 0:
                continue
 
            #유효성 확인, valid가 0이면 복사한 데이터라는 의미임
            if anno['valid'] == 0:
                continue
        
            label_dict[cat].append(anno['action_id'][cat])
    
    #카테고리별 데이터 생성
    ret_label = dict()
    for cat in category:
        #유효성 체크, 클래스 리스트에 데이터가 들어있지 않으면 -1 설정
        if len(label_dict[cat]) == 0:
            ret_label['gt_'+cat] = -1
            continue
 
        #빈도수 확인
        values, counts = np.unique(np.array(label_dict[cat]), return_counts=True)
        max_count = np.max(counts) #최대 빈도수
        labels = values[counts == max_count] #최대 빈도수 클래스
        labels = labels.tolist()
        ret_label['gt_'+cat] = labels[0]
 
    return ret_label
 
def make_image_feat_data(anno_list, img_feats):
    """
    학습을 위한 이미지 특징을 생성하는 기능
    args:
        anno_list (list[dict]): 어노테이션 데이터 리스트, 각 항목은 'image_id'와 'valid' 키를 포함하고 있어야함
        img_feats (dict): 이미지 아이디에 대응하는 이미지 특징, 이미지 특징은 Tenosr 타입이고 shpae은 (C,)임
    return (Tensor): 데이터에 대한 이미지 특징, shape (T, C, 1)
    """
 
    #이미지 특징 크기를 미리 셋팅해둠
    feat_size = 0 #이미지 특징 크기
    for anno in anno_list:
        image_id = anno['image_id'] #이미지 아이디
 
        #유효성 체크
        if anno['valid'] == 0: 
            continue
        
        #이미지 특징 크기 확인
        feat_size = img_feats[image_id].shape[0]
        break
 
    #여러 프레임의 이미지 특징을 하나의 변수로 만듦
    img_feat_list = [] #이미지 특징을 저장할 리스트
    for anno in anno_list:
        image_id = anno['image_id'] #이미지 아이디
 
        #유효성 체크
        if anno['valid'] == 0:
            img_feat_list.append(torch.zeros(feat_size)) #0값으로 설정함
        else:
            img_feat_list.append(img_feats[image_id]) #이미지 특징 저장
 
    #여러프레임의 이미지 특징을 반환
    return torch.stack(img_feat_list, dim=1).unsqueeze(-1) 
    
def gen_pkl_data_v22(anno_list, train_category, category, target_label, category_margin_ratio, h, w, image_folder, img_feats=None):
    """
    학습을 위한 데이터 1개를 생성하는 기능
    학습 카테고리별로 학습 클래스를 설정함
    category(메인 카테고리)의 클래스는 데이터의 중심 40% 구간에 target_label이 존재하면 target_label로 설정됨, target_label이 존재하지 않는다면 데이터는 생성하지 않음
    
    args:
        anno_list (list[dict]): 어노테이션 데이터 리스트
        train_category (list[str]): 학습 카테고리 리스트
        category (str): 학습 데이터의 메인 카테고리
        target_label (int): 메인 카테고리에 해당하는 타겟 클래스, target_label 정보가 존재해야 학습 데이터를 생성함
        category_margin_ratio (dict): 카테고리별 클래스 결정을 위한 데이터의 시작 위치(%), 클래스를 결정하는데 사용하는 전체 범위는 데이터 중앙의 (1-2*ratio) *100 %가 됨
        h (int): 원본 영상의 높이
        w (int): 원본 영상의 너비
        image_folder (str): 영상 폴더
        img_feats (dict): 이미지 아이디에 대응하는 이미지 특징, 이미지 특징은 Tenosr 타입이고 shpae은 (C,)임, 파라미터를 설정하지 않으면 이미지 특징은 학습데이터에 포함하지 않음
    return (dict): 학습데이터
        keypoint (np.array): 스켈레톤 좌표
        keypoint_score (np.array): 스켈레톤 스코어
        frame_dir (str): 영상 폴더
        img_shape (tuple): 원본 영상의 크기
        original_shape (tuple): 원본 영상의 크기
        total_frames (int): 데이터의 프레임 수
        category (str): 샘플의 메인 카테고리
        gt_{category_key} (int): 카테고리별 클래스
        image_feature (Tensor): 이미지 특징, shape (T, C, 1)
    """
 
    #유효성 체크
    if check_invaild(anno_list, 0.5):
        return None
 
    
    label = dict()
    #메인 카테고리의 클래스 설정
    if check_label_v22(anno_list, category, target_label, 0.3):
        label['gt_'+category] = target_label
    else:
        return None
    
    #메인 카테고리를 제외한 나머지 카테고리의 클래스 설정
    non_target_category = remove_items(train_category, [category]) #메인 카테고리를 제외한 나머지 카테고리
    label.update(make_label_v22(anno_list, non_target_category, category_margin_ratio))
    
    #학습 데이터 생성
    pkl_data = dict()
    
    kps = []
    for anno in anno_list:
        kps.append(anno['keypoints'])
        
    kps = np.array([kps])    
    pkl_data['keypoint'] = kps[..., :2]
    pkl_data['keypoint_score'] = kps[..., 2]
    pkl_data['frame_dir'] = image_folder
    pkl_data['img_shape'] = (h, w)
    pkl_data['original_shape'] = (h, w)
    pkl_data['total_frames'] = kps.shape[1]
    pkl_data['category'] = category
    pkl_data.update(label)
 
    #이미지 특징 설정
    if img_feats is not None:
        pkl_data['image_feature'] = make_image_feat_data(anno_list, img_feats)
 
    return pkl_data
 
def make_action_pkl_file_v22(data_manager, train_category, category, target_label, category_margin_ratio, k, img_feat_dict=None):
    """
    하나의 클립에서 학습 데이터를 생성하는 기능
    학습 카테고리별로 학습 클래스를 설정함
    args:
        data_manager: AnnotationDataManager 객체, 어노테이션 데이터파일로 초기화까지 진행
        train_category (list[str]): 학습 카테고리 리스트
        category (str): 학습 데이터의 메인 카테고리
        target_label (int): 메인 카테고리에 해당하는 타겟 클래스, target_label 정보가 존재해야 학습 데이터를 생성함
        category_margin_ratio (dict): 카테고리별 클래스 결정을 위한 데이터의 시작 위치(%), 클래스를 결정하는데 사용하는 전체 범위는 데이터 중앙의 (1-2*ratio) *100 %가 됨
        k (int): 하나의 데이터를 만드는데 사용하는 프레임 숫자
        img_feat_dict (dict): 추적 아이디의 이미지 특징 데이터, 이미지 특징 데이터는 dict 구조로 이미지 아이디별 각 프레임에서의 이미지 특징을 담고 있음.
    return (dict): 하나의 클립에 대한 학습 데이터 리스트, 학습 데이터는 추적 아이디별로 저장됨, 추적 아이디의 학습 데이터는 list[dict]의 구조임
    """
 
    
    
    #이미지 크기 및 폴더 셋팅
    # h, w = data_manager.read_image().shape[:2]
    h = 1080
    w = 1920
    image_folder = data_manager.get_image_folder()
    
    #클립의 시작, 끝 프레임 설정
    start_frame = data_manager.dataset['info']['start']
    end_frame = data_manager.dataset['info']['end']
 
    #추적 아이디를 기준으로 데이터 형태 변환
    track_data = data_manager.get_track_anno_data(filter=True, start_frame=start_frame, end_frame=end_frame+1)
 
    #학습 데이터 생성, 추적 아이디별로 생성함
    pkl_datas = {}
    for track_id, anno_list in track_data.items():
 
        #이미지 특징이 있을경우 셋팅함
        img_feats = None
        if img_feat_dict is not None:
            img_feats = img_feat_dict[track_id]
 
        frame_no = len(anno_list)
 
        image_idx_1st = data_manager.image_ids.index(anno_list[0]['image_id']) #시작 프레임
        image_idx_end = data_manager.image_ids.index(anno_list[-1]['image_id']) #끝 프레임
 
        #유효성 체크, 데이터의 길이가 설정한 학습데이터의 길이보다 작을 경우 사용하지 않음
        interval = image_idx_end - image_idx_1st
        if interval < k:
            continue
        
        #학습 클래스 변환, 라벨링된 데이터를 학습을 위한 데이터로 변환 
        simple_anno_list = cvt_labelmap_v22(anno_list)
 
        #유효성 셋팅
        for i in range(0, frame_no):
            simple_anno_list[i]['valid'] = 1
 
        #추적 아이디의 데이터를 저장할 리스트 생성
        pkl_datas[track_id] = []
 
        #pkl data 생성
        anno_q = [simple_anno_list[0]]
        for i in range(1, frame_no):
 
            #인접한 데이터의 프레임 확인
            image_idx_end = data_manager.image_ids.index(simple_anno_list[i]['image_id'])
            image_idx_end_m1 = data_manager.image_ids.index(anno_q[-1]['image_id'])
 
            #인접한 데이터의 프레임 간격에 따라서 데이터 처리
            interval = image_idx_end - image_idx_end_m1
            if interval > k: #k보다 클경우 다른 클립으로 분리함
                anno_q = []
            elif interval > 1: #k보다 작고 1보다 클경우 더미 데이터를 넣어줌
                anno_q.extend(make_dummy_list(anno_q[-1], interval-1))
 
            #큐에 데이터를 붙여줌
            anno_q.append(simple_anno_list[i])
            
            #큐의 데이터가 k개 이상인 경우만 학습 데이터 생성 프로세스를 진행함
            frame_num = len(anno_q)
            if frame_num < k:
                continue
 
            #학습 데이터 생성
            pkl_data = gen_pkl_data_v22(anno_q[frame_num-k:], train_category, category, target_label, category_margin_ratio, h, w, image_folder, img_feats=img_feats)
 
            #학습 데이터가 생성되지 않았을 경우
            if pkl_data is None:
                continue
            
            #추적 아이디의 키에 피클 데이터 저장
            pkl_datas[track_id].append(pkl_data)
        
        #추적 아이디의 학습 데이터가 없을 경우 키 제거
        if len(pkl_datas[track_id]) == 0:
            del pkl_datas[track_id]
 
    return pkl_datas
 
 
def run_generate_pkl_data_from_json_v22(args):
    """
    클립 정보를 담고있는 어노테이션 파일들을 사용해서 학습데이터를 만드는 기능
    args:
        category_k (dict): 카테고리별 학습데이터를 생성하는 프레임 수
        train_category (list[str]): 학습 카테고리 리스트
        category_margin_ratio (dict): 카테고리별 클래스 결정을 위한 데이터의 시작 위치(%), 클래스를 결정하는데 사용하는 전체 범위는 데이터 중앙의 (1-2*ratio) *100 %가 됨
        data_root (str): 클립 정보를 담고있는 어노테이션 파일들이 저장된 폴더
        pkl_root (str): 학습데이터를 저장할 폴더
        image_feat_folder (str): 이미지 특징이 저장된 폴더
    """
 
    #어노테이션 관리자 설정
    data_manager = AnnotationDataManager()
 
    #학습데이터 생성을 위한 카테고리별 프레임 수
    category_k = args.category_k
    train_category = args.train_category #학습 카테고리 리스트
    category_margin_ratio = args.category_margin_ratio #카테고리별 클래스 설정 구간
 
    data_root = args.data_root #클립 저장 폴다
    for category, k in category_k.items():
        #클립 리스트
        dataset_folder = os.path.join(data_root, category)
        json_name_list, json_path_list = search_file(dataset_folder, '.json')
 
        #카테고리 별 저장 폴더 생성
        save_folder = os.path.join(args.pkl_root, category)
        for i, (json_path, json_name) in enumerate(zip(json_path_list, json_name_list)):
            target_label = json_path.split('/')[-2]
            cat_label = os.path.join(save_folder, target_label)
            create_directory(cat_label)
 
        file_num = len(json_name_list) #클립 수량, 콘솔 출력용
        total_data = 0 #학습 데이터 수량, 콘솔 출력용
        for i, (json_path, json_name) in enumerate(zip(json_path_list, json_name_list)):
            # print(f"{category} -> [ {i+1} / {file_num} ] {json_name}")
 
            #어노테이션 로드
            data_manager.load_annotation(json_path)
            data_manager.init_annotation()
 
            #타겟 라벨 -> 메인 카테고리의 클래스
            target_label = int(json_path.split('/')[-2])
 
            #영상 특징 사용 여부
            if args.image_feat_folder is not None:
                img_feat_file = json_name[:17] + '.pkl'
                img_feat_path = os.path.join(args.image_feat_folder, img_feat_file)
                img_feat_dict = load_pkl_data(img_feat_path)
 
            #피클 데이터 생성
            pkl_datas = make_action_pkl_file_v22(data_manager, train_category, category, target_label, category_margin_ratio, k, img_feat_dict=img_feat_dict)
            
            #학습 데이터 저장
            for track_id, datas in pkl_datas.items():
                pkl_name = json_name.replace('json', 'pkl')
                save_pkl_path = os.path.join(save_folder, f'{target_label:02d}', pkl_name)
                save_pkl_data(datas, save_pkl_path)
                total_data += len(datas)
                print(f"{category} -> [ {i+1} / {file_num} ] {save_pkl_path} : {len(datas):5d}/{total_data:}")