action_dataset_loader.py

Go to the documentation of this file.

 
import torch
import torch.nn as nn
from smutils.utils_os import search_file
from smutils.utils_data import load_pkl_data
import random
 
import copy
import numpy as np
import os
from smaction.datasets.transform.pose_transform import *
 
#피클 데이터의 구조
#    dict 구조
#        keys 필수 -> 'label'
#        행동 인식을 위한 key 예시 -> 'keypoints'
 
class ActionDatasetLoader(torch.utils.data.Dataset):
    """
    세이프모션의 학습용 데이터를 기반으로하는 데이터로더
    학습용 데이터는 '.pkl'파일로 구성됨
    TODO: 사용하지 않을 가능성이 커보임, 삭제 요망
    args:
        mode (str): 데이터로더의 셋팅 모드, 학습용일 경우 'train', 평가일 경우 'val'
        train_data_folder (str): 학습용 데이터가 저장된 폴더 또는 파일
        test_data_folder (str): 평가용 데이터가 저장된 폴더 또는 파일
        clip_len (int): 샘플링할 프레임 수
    """
    def __init__(self, mode, train_data_folder, test_data_folder, clip_len):
        self.mode = mode
        self.train_data_folder = train_data_folder
        self.test_data_folder = test_data_folder
        self.clip_len = clip_len
 
        self.load_data()
        
    def load_data(self):
        """
        데이터를 로드하는 기능
        """
 
        #모드에 따른 데이터 경로 설정
        data_folder = self.train_data_folder if self.mode == 'train' else self.test_data_folder
        if os.path.isdir(data_folder): #데이터 경로가 폴더일 경우
            #폴더 내부의 .pkl 파일을 모두 검색
            name_list, path_list = search_file(data_folder, '.pkl')
 
            #데이터 로드
            self.dataset = []
            for name, path in zip(name_list, path_list):
                pkl = load_pkl_data(path)
                self.dataset.append(pkl)
        else:
            #데이터 로드
            self.dataset = load_pkl_data(data_folder)
    
    def preprocessing(self, sample):
        """
        모델의 입력을 위한 데이터 전처리 기능
        스켈레톤의 히트맵을 생성하는게 메인 기능임, 데이터 증강(data augmentation) 포함됨
        TODO: 데이터 증강은 학습 모드일 경우만 진행되도록 코드 변경 필요
        args:
            sample (dict): 모델의 입력을 생성하기 위한 데이터
        return (dict): 모델의 입력이 포함된 데이터
        """
 
        sample = pose_sampling(sample, clip_len=self.clip_len) #샘플링
        sample = pose_shift(sample, shift_ratio=0.05) #스켈레톤 이동, 데이터 증강
        sample = pose_compact(sample) #스켈레톤 범위 셋팅
        sample = pose_resize(sample, scale=(64, 64)) #리사이즈
        sample = pose_random_crop(sample, area_range=(0.56, 1.0), aspect_ratio_range=(3 / 4, 4 / 3)) #랜덤 크롭, 데이터 증강
        sample = pose_resize(sample, scale=(56, 56)) #리사이즈, 크롭을 진행했기 때문에 다시 리사이즈 해줌
        sample = pose_flip(sample, flip_ratio=0.5) #좌우반전, 데이터 증강
        sample = make_pose_heatmap(sample) #히트맵 생성
        
        return sample
 
    def __len__(self):
        return len(self.dataset)
 
    def __getitem__(self, index):
        sample = self.dataset[index]
        if isinstance(sample, list): #샘플이 리스트일 경우 랜덤으로 한개 선택함
            size_data = len(sample) #리스트 크기
 
            #랜덤으로 선택
            rand_idx = random.randint(0, size_data-1)
            sample = sample[rand_idx]
        
        return self.preprocessing(sample)
 
class ActionDatasetLoader_v22(torch.utils.data.Dataset):
    """
    세이프모션의 데이터셋을 기반으로 단일 경로 모델을 학습하기 위한 데이터로더
    args:
        mode (str): 데이터로더의 셋팅 모드, 학습용일 경우 'train', 평가일 경우 'val'
        data_folder (str): 학습용 데이터셋이 저장된 폴더, 해당 폴더 안에는 mode 폴더가 포함되어 있어야함
        category_info (dict): 학습 데이터의 카테고리별 클래스 수
            TODO: category_info는 제거해도 동작가능할 것으로 보여짐, 코드 수정 필요
        clip_len (int): 샘플링할 프레임 수
    """
    def __init__(self, mode, data_folder, category_info, clip_len):
        self.mode = mode
        self.data_folder = data_folder
        self.category_info = category_info
        self.clip_len = clip_len
 
        self.load_data()
        
    def load_data(self):
        """
        데이터를 로드하는 기능
        """
        self.dataset = []
        for category, num in self.category_info.items(): #카테고리, 클래스 수
            folder = os.path.join(self.data_folder, self.mode, category) #카테고리-클래스의 데이터 폴더
            name_list, path_list = search_file(folder, '.pkl') #해당 폴더의 모든 pkl파일 검색
            
            #데이터 로드
            data_num = len(path_list)
            for i, (name, path) in enumerate(zip(name_list, path_list)):
                print(f'{self.mode} : {category}  [{i+1}/{data_num}] -> {name}', end='\r')
                pkl = load_pkl_data(path)
                self.dataset.append(pkl)
            
    def preprocessing(self, sample):
        """
        모델의 입력을 위한 데이터 전처리 기능
        스켈레톤의 히트맵을 생성하는게 메인 기능임, 데이터 증강(data augmentation) 포함됨
        TODO: 데이터 증강은 학습 모드일 경우만 진행되도록 코드 변경 필요
        args:
            sample (dict): 모델의 입력을 생성하기 위한 데이터
        return (dict): 모델의 입력이 포함된 데이터
        """
        sample = pose_sampling(sample, clip_len=self.clip_len) #샘플링
        sample = pose_shift(sample, shift_ratio=0.05) #스켈레톤 이동, 데이터 증강
        sample = pose_compact(sample) #스켈레톤 범위 셋팅
        sample = pose_resize(sample, scale=(64, 64)) #리사이즈
        sample = pose_random_crop(sample, area_range=(0.56, 1.0), aspect_ratio_range=(3 / 4, 4 / 3)) #랜덤 크롭, 데이터 증강
        sample = pose_resize(sample, scale=(56, 56)) #리사이즈, 크롭을 진행했기 때문에 다시 리사이즈 해줌
        sample = pose_flip(sample, flip_ratio=0.5) #좌우 반전, 데이터 증강
        sample = make_pose_heatmap(sample) #히트맵 생성
        
        return sample
 
    def __len__(self):
        return len(self.dataset)
 
    def __getitem__(self, index):
        sample = self.dataset[index]
        if isinstance(sample, list): #샘플이 리스트일 경우 한개 선택함
            size_data = len(sample)
 
            #인덱스 설정
            if self.mode == 'train': #학습 모드일 경우 랜덤으로 선택
                rand_idx = random.randint(0, size_data-1)
            else: #평가모드일 경우 중심에서 한개 선택
                rand_idx = size_data//2
            
            sample = sample[rand_idx]
        
        return self.preprocessing(sample)
 
 
 
class ActionDatasetLoader_mtml(torch.utils.data.Dataset):
    """
    세이프모션의 데이터셋을 기반으로 2-스트림 경로 모델을 학습하기 위한 데이터로더
    args:
        mode (str): 데이터로더의 셋팅 모드, 학습용일 경우 'train', 평가일 경우 'val'
        data_folder (str): 학습용 데이터셋이 저장된 폴더, 해당 폴더 안에는 mode 폴더가 포함되어 있어야함
        category_info (dict): 학습 데이터의 카테고리별 클래스 수
            TODO: category_info는 제거해도 동작가능할 것으로 보여짐, 코드 수정 필요
        clip_len_action (int): 행동 경로에 사용되는 프레임 수
        clip_len_pose (int): 포즈 경로에 사용되는 프레임 수
    """
    def __init__(self, mode, data_folder, category_info, clip_len_action=20, clip_len_pose=8):
        self.mode = mode
        self.data_folder = data_folder
        self.category_info = category_info
        self.clip_len_action = clip_len_action
        self.clip_len_pose = clip_len_pose
 
        self.load_data()
        
    def load_data(self):
        """
        데이터를 로드하는 기능
        """
        self.dataset = []
        for category, num in self.category_info.items(): #카테고리, 클래스 수
            folder = os.path.join(self.data_folder, self.mode, category) #카테고리-클래스의 데이터 폴더
            name_list, path_list = search_file(folder, '.pkl') #해당 폴더의 모든 pkl파일 검색
            
            #데이터 로드
            data_num = len(path_list)
            for i, (name, path) in enumerate(zip(name_list, path_list)):
                print(f'{self.mode}/{category}/{num:02d}  [{i+1}/{data_num}] -> {name}', end='\r')
                pkl = load_pkl_data(path)
                self.dataset.append(pkl)
            
    def preprocessing(self, sample):
        """
        모델의 입력을 위한 데이터 전처리 기능
        스켈레톤의 히트맵을 생성하는게 메인 기능임, 데이터 증강(data augmentation) 포함됨
        TODO: 데이터 증강은 학습 모드일 경우만 진행되도록 코드 변경 필요
        args:
            sample (dict): 모델의 입력을 생성하기 위한 데이터
        return (dict): 모델의 입력이 포함된 데이터
        """
 
        sample = pose_sampling(sample, clip_len=self.clip_len_action) #샘플링
        sample = pose_shift(sample, shift_ratio=0.01) #스켈레톤 이동, 데이터 증강
        sample = pose_compact(sample) #스켈레톤 범위 셋팅
        sample = pose_resize(sample, scale=(64, 64)) #리사이즈
        sample = pose_random_crop(sample, area_range=(0.56, 1.0), aspect_ratio_range=(3 / 4, 4 / 3)) #랜덤 크롭, 데이터 증강
        sample = pose_resize(sample, scale=(56, 56)) #리사이즈, 크롭을 진행했기 때문에 다시 리사이즈 해줌
        sample = pose_flip(sample, flip_ratio=0.5) #좌우 반전, 데이터 증강
 
        sample_pose = copy.deepcopy(sample) #포즈 경로의 입력 데이터 생성을 위해 복사
 
        #포즈 경로의 입력 범위 샘플링
        frames = sample_pose['keypoint'].shape[1]
        start = (frames - self.clip_len_pose)//2
        end = start+self.clip_len_pose
        sample_pose['keypoint'] = sample_pose['keypoint'][:,start:end, :, : ]
        sample_pose['keypoint_score'] = sample_pose['keypoint_score'][:,start:end, : ]
 
        #포즈 경로의 입력 데이터 생성을위한 스켈레톤 좌표 조정
        sample_pose = pose_compact(sample_pose) #스켈레톤 범위 셋팅
        sample_pose = pose_resize(sample_pose, scale=(56, 56)) #리사이즈
 
        #히트맵 생성
        sample = make_pose_heatmap(sample)
        sample_pose = make_pose_heatmap(sample_pose)
        sample['pose_heatmap_for_action'] = sample['pose_heatmap']
        sample['pose_heatmap_for_pose'] = sample_pose['pose_heatmap']
        
        return sample
 
    def __len__(self):
        return len(self.dataset)
 
    def __getitem__(self, index):
        sample = self.dataset[index]
        if isinstance(sample, list): #샘플이 리스트일 경우 한개 선택함
            size_data = len(sample)
 
            #인덱스 설정
            if self.mode == 'train': #학습 모드일 경우 랜덤으로 선택
                rand_idx = random.randint(0, size_data-1)
            else: #평가모드일 경우 중심에서 한개 선택
                rand_idx = size_data//2
            
            sample = sample[rand_idx]
        
        return self.preprocessing(sample)