utils_data.py

Go to the documentation of this file.

import random
import copy
import pickle
import shutil
import os
from smutils.utils_os import search_file, create_directory
 
#TODO: utils_data.py, utils_os.py 에 들어갈 함수목록 체크
#TODO: utils_data.py에서 smdataset으로 이동할지 결정
 
def load_pkl_data(pkl_path):
    """
    pkl 데이터를 로드하는 기능
    args:
        pkl_path (str): pkl 데이터가 저장된 경로
    return: pkl 데이터
    """
    with open(pkl_path, 'rb') as f:
        pkl_datas = pickle.load(f)
    return pkl_datas
 
def save_pkl_data(pkl_data, save_path):
    """
    pkl 데이터를 저장하는 기능
    args:
        pkl_data : 저장할 데이터
        save_path (str): pkl 데이터를 저장할 경로
    """
    with open(save_path, 'wb') as f:  
        pickle.dump(pkl_data, f)
 
def dataset_class_filtering(pkl_datas, active_class):
    """
    특정 클래스의 데이터만 남기고 제거하는 기능
    args:
        pkl_datas (list[dict] or str): 입력 데이터 리스트 또는 데이터의 경로, 데이터에는 'label' 항목이 포함되어 있어야함
        active_class (list) : 남기고 싶은 데이터
    return (list[dict]): 필터링된 데이터
    """
 
    if isinstance(pkl_datas, str): #pkl_datas가 문자열이면 데이터 로드
        with open(pkl_datas, 'rb') as f:
            pkl_datas = pickle.load(f)
 
    #데이터 필터링
    filtered_datas = []
    for data in pkl_datas:
        if data['label'] in active_class: #남길 데이터인지 확인
            filtered_datas.append(data)
    return filtered_datas
 
def print_data_num_per_class(data_list, labelmap):
    """
    클래스별 데이터 수량을 출력하는 기능
    args:
        data_list (list[dict]): 데이터 리스트, 데이터에는 클래스 정보를 담고있는 'label'을 포함하고 있어야함
        labelmap (dict): 클래스 정보를 담고있는 변수, key는 클래스 번호이고 value는 클래스 이름임
    """
    if isinstance(data_list, str):
        with open(data_list, 'rb') as f:
            data_list = pickle.load(f)
            
    bins_data = dict()
    for i in range(len(labelmap.keys())):
        bins_data[i] = 0
        
    for data in data_list:    
        bins_data[data['label']] += 1
    
    for label, cnt in bins_data.items():
        print(f"{labelmap[label]} : {cnt}")
    
 
def split_train_and_val(datas, class_num, max_train_data_num, max_val_data_num):
    """
    데이터를 훈련용과 평가용으로 분리하는 기능
    args:
        datas (list[dict]): 데이터 리스트
        class_num (int): 총 클래스 수
        max_train_data_num (int): 훈련 데이터셋 최대 수량
        max_val_data_num (int): 평가 데이터셋 최대 수량
    return:
        train_list (list[dict]): 훈련용 데이터셋
        val_list (list[dict]): 평가용 데이터셋
    """
 
    #클래스별 데이터 분리
    data_dict = dict()
    for i in range(class_num):
        data_dict[i] = [] #초기화
    
    #분리
    for data in datas:
        data_dict[data['label']].append(data)    
    
    #훈련 및 평가용 데이터 분리
    train_list = []
    val_list = []
    for key, data_list in data_dict.items(): #클래스별 데이터 리스트
        data_num = len(data_list) #데이터 수량
        train_data_num = max_train_data_num #훈련용 데이터 수량
        if data_num < max_train_data_num: #데이터 수량이 부족할 경우, 95%를 학습용으로 사용
            train_data_num = int(data_num * 0.95)
        
        #평가용 데이터 끝 인덱스
        val_data_idx = min(data_num, train_data_num+max_val_data_num)
        
        #분리
        random.shuffle(data_list) #섞기
        train_list.extend(copy.deepcopy(data_list[:train_data_num])) #훈련용 데이터
        val_list.extend(copy.deepcopy(data_list[train_data_num:val_data_idx])) #평가용 데이터
    
    #섞기
    random.shuffle(train_list)
    random.shuffle(val_list)
    
    return train_list, val_list
 
def remove_items(original_list, remove_items):
    """
    리스트에서 특정 원소를 제거하는 기능
    args:
        original_list (list): 입력 데이터
        remove_items (list): 제거하려는 원소
    return (list): 특정 원소가 제거된 리스트
    """
    return [item for item in original_list if item not in remove_items]
    
def make_pkl_info(pkl_name_list, pkl_path_list, key='label'):
    """
    피클 파일의 정보를 생성, 입력한 키의 원소별 파일 수량과 피클 파일에 포함된 키의 원소를 생성하는 기능
    args:
        pkl_name_list (list[str]): 피클 데이터 파일의 이름 리스트
        pkl_path_list (list[str]): 피클 데이터 파일 경로 리스트
    """
    data_info_dict = dict()
    file_info_dict = dict()
    for name, path in zip(pkl_name_list, pkl_path_list):
        pkl_list = load_pkl_data(path) #피클 파일 로드
        file_info_dict[name] = [] #파일에 포함된 키 데이터 초기화
 
        #정보 생성
        for pkl in pkl_list:
            item = pkl[key] #키의 원소
 
            #원소별 파일 초기화
            if item not in data_info_dict:
                data_info_dict[item] = []
            
            #파일에 포함된 원소 추가
            if item not in file_info_dict[name]:
                file_info_dict[name].append(item)
            
            #원소별 파일 추가
            if name not in data_info_dict[item]:
                data_info_dict[item].append(name)
                
    #정렬
    data_info_dict = dict(sorted(data_info_dict.items()))
    return data_info_dict, file_info_dict
 
 
# def select_items(pkl_folder, ratio=0.2):
#     #하나에 클립에 여러개의 라벨이 있을 수 있음
#     #최소한의 클립으로 여러 라벨이 골고루 커버가능하도록 함
#     pkl_name_list, pkl_path_list = search_file(pkl_folder, '.pkl')
#     label_info, file_info = make_pkl_info(pkl_name_list, pkl_path_list)
 
#     # 클립 수로 오름차순으로 정렬
#     sorted_data = copy.deepcopy(dict(sorted(label_info.items(), key=lambda item: len(item[1]))))
#     item_dict = {}
#     select_item_list = []
#     for key in label_info.keys():
#         item_dict[key] = []
    
#     while True:
#         #클립 수가 가장 적은 라벨 선택
#         label = list(sorted_data.keys())[0]
#         clip_list = sorted_data[label]
 
#         #선택할 클립의 수 결정
#         sample_size = max(1, int(len(clip_list)*ratio)) # 선택해야하는 클립 수
#         sample_size = max(0, sample_size - len(item_dict[label])) # 선택해야하는 클립 수에 이미 선택된 클립수를 뺌
        
#         #클립 선택
#         selected_clip = random.sample(clip_list, sample_size)
#         select_item_list.extend(selected_clip)
#         for clip in selected_clip:
#             l_list = file_info[clip]
#             for l in l_list:
#                 item_dict[l].append(clip)
        
#         #선택된 클립 제거
#         for key, value_list in sorted_data.items():
#             sorted_data[key] = [val for val in value_list if val not in selected_clip]
        
#         del sorted_data[label]
        
#         if not sorted_data:
#             break
        
#         # 클립 수로 오름차순으로 정렬
#         sorted_data = dict(sorted(sorted_data.items(), key=lambda item: len(item[1])))
    
#     remain_item_list = remove_items(pkl_name_list, select_item_list)
 
#     return select_item_list, remain_item_list
 
 
def split_train_and_val_v22(data_folder, save_folder, category_info, total_data_num=100, ratio=0.8):
    """
    세이프모션의 2.2 버전 행동 데이터를 훈련용과 평가용으로 분리하는 기능
    args:
        data_folder (str): 데이터가 저장된 폴더
        save_folder (str): 학습 및 훈련용 데이터를 저장할 폴더
        category_info (dict): 학습 데이터의 카테고리별 클래스 수
        total_data_num (int): 클래스별 총 데이터 수(학습+평가 데이터 수)
        ratio (float): 학습용 데이터의 비율
    """
    #변수 초기화 및 저장 폴더 생성
    split_datas = dict() #클래스별 데이터 분리용 변수
    for category, num in category_info.items():
        split_datas[category] = dict()
        for label in range(num):
            #클래스별 데이터 분리용 변수 초기화
            split_datas[category][label] = dict()
            split_datas[category][label]['train'] = []
            split_datas[category][label]['val'] = []
            
            #저장 폴더 생성
            folder_train = os.path.join(save_folder, 'train', category, f'{label:02d}') #훈련 데이터 폴더
            folder_val = os.path.join(save_folder, 'val', category, f'{label:02d}') #평가 데이터 폴더
            create_directory(folder_train)
            create_directory(folder_val)
 
    #데이터 분리
    for category, num in category_info.items():#카테고리별 클래스 수량
        for label in range(num): #클래스 번호
 
            folder = os.path.join(data_folder, category, f'{label:02d}') #클래스 관련 데이터 폴더
            name_list, path_list = search_file(folder, '.pkl') #폴더 내 데이터 목록
 
            data_num = min(len(path_list), total_data_num)  #총 데이터수
            train_num = int( data_num * ratio ) #학습용 데이터 수 설정
            
            random.shuffle(path_list) #섞기
            split_datas[category][label]['train'] = path_list[:train_num] #학습용 데이터
            split_datas[category][label]['val'] = path_list[train_num:data_num] #평가용 데이터
    
    #데이터 저장
    for category, num in category_info.items():
        for label in range(num):
            for mode, path_list in split_datas[category][label].items():
                folder = os.path.join(save_folder, mode, category, f'{label:02d}')
                for path in path_list:
                    name = path.split('/')[-1]
                    save_path = os.path.join(folder, name)
                    shutil.copy(path, save_path)
 
def load_labelmap(path):
    """
    레이블 맵을 로드하는 기능
    args:
        path (str): 레이블 맵이 저장된 경로, 레이블 맵 파일은 1라인에 클래스 명이 기록되어 있어야함
    return (list[str]): 레이블 맵
    """
    label_map = [x.strip() for x in open(path).readlines()]
    return label_map