add node mask and action selector

README.md

@@ -9,10 +9,10 @@ An **unofficial** PyTorch implementation of the paper "Skeleton-Based Action Rec
 
 ## Dependencies
 
-- Python >= 3.5
+- Python >= 3.6
 - scipy >= 1.3.0
 - numpy >= 1.16.4
-- PyTorch >= 1.1.0
+- PyTorch >= 1.4.0
 - tensorboardX >= 1.8   (For logging)
 
 ## Directory Structure
@@ -35,22 +35,21 @@ Most of the interesting stuff can be found in:
 3. Generate the joint dataset first:
 
 ```bash
-cd data_gen
-python3 ntu_gen_joint_data.py
+python3 ./data_gen/ntu_gen_joint_data.py
 ```
 
 Specify the data location if the raw skeletons data are placed somewhere else. The default looks at `./data/nturgbd_raw/`.
 
 4. Then, in `data_gen/`, generate the bone dataset:
 
 ```bash
-python3 ntu_gen_bone_data.py
+python3 ./data_gen/ntu_gen_bone_data.py
 ```
 
 5. Finally, generate the motion data from joints/bones:
 
 ```bash
-python3 ntu_gen_motion_data.py
+python3 ./data_gen/ntu_gen_motion_data.py
 ```
 
 The generation scripts look for generated data in previous step. By default they look at `./data`; change dir configs if needed.

config/node_config.yaml

@@ -0,0 +1,4 @@
+# node_mask: [1, 2, 5, 9, 13, 14, 17, 18, 21] # quarter joint
+# node_mask: [1, 2, 3, 4, 5, 6, 9, 10, 13, 14, 17, 18, 21] # harf joint
+# selected_action: ['A004', 'A005', 'A006', 'A007', 'A008', 'A009', 'A010', 'A027', 'A031', 'A059']
+selected_action: ['A004', 'A005']

data_gen/ntu_gen_bone_data.py

@@ -4,52 +4,62 @@
 
 from tqdm import tqdm
 
+# Joint 연관 관계
 paris = {
-    'ntu/xview': (
+    'ntu/xview': [
         (1, 2), (2, 21), (3, 21), (4, 3), (5, 21),
         (6, 5), (7, 6), (8, 7), (9, 21), (10, 9),
         (11, 10), (12, 11), (13, 1), (14, 13), (15, 14),
         (16, 15), (17, 1), (18, 17), (19, 18), (20, 19),
         (22, 23), (21, 21), (23, 8), (24, 25), (25, 12)
-    ),
-    'ntu/xsub': (
+    ],
+
+    'ntu/xsub': [
         (1, 2), (2, 21), (3, 21), (4, 3), (5, 21),
         (6, 5), (7, 6), (8, 7), (9, 21), (10, 9),
         (11, 10), (12, 11), (13, 1), (14, 13), (15, 14),
         (16, 15), (17, 1), (18, 17), (19, 18), (20, 19),
         (22, 23), (21, 21), (23, 8), (24, 25), (25, 12)
-    ),
+    ],
 
-    'kinetics': (
+    'kinetics': [
         (0, 0), (1, 0), (2, 1), (3, 2), (4, 3), (5, 1),
         (6, 5), (7, 6), (8, 2), (9, 8), (10, 9), (11, 5),
         (12, 11), (13, 12), (14, 0), (15, 0), (16, 14), (17, 15)
-    )
+    ]
 }
 
 sets = {'train', 'val'}
 datasets = {'ntu/xview', 'ntu/xsub'}
 
 def gen_bone_data():
     """Generate bone data from joint data for NTU skeleton dataset"""
-    for dataset in datasets:
-        for set in sets:
+    for dataset in datasets: # 'ntu/xview', 'ntu/xsub'
+        for set in sets: # 'train', 'val'
             print(dataset, set)
-            data = np.load('../data/{}/{}_data_joint.npy'.format(dataset, set))
-            N, C, T, V, M = data.shape
+
+            # 앞에서 생성한 Joint data Load
+            data = np.load('./data/{}/{}_data_joint.npy'.format(dataset, set))
+            N, C, T, V, M = data.shape # channels (C), frames (T), nodes (V), persons (M)
+
+            # Bone 데이터 저장할 파일 생성. Joint 데이터와 Shape를 같게.
             fp_sp = open_memmap(
-                '../data/{}/{}_data_bone.npy'.format(dataset, set),
+                './data/{}/{}_data_bone.npy'.format(dataset, set),
                 dtype='float32',
                 mode='w+',
                 shape=(N, 3, T, V, M))
-
-            # Copy the joints data to bone placeholder tensor
-            fp_sp[:, :C, :, :, :] = data
-            for v1, v2 in tqdm(paris[dataset]):
+
+            # Joint data를 fp_sp에 복사 
+            fp_sp[:, :C, :, :, :] = data # Deep copy
+
+            for v1, v2 in tqdm(paris[dataset]): # dataset : 'ntu/xview', 'ntu/xsub'
                 # Reduce class index for NTU datasets
+
+                # Joint 번호는 1부터 시작하지만, 인덱스는 0부터 시작하기 때문
                 if dataset != 'kinetics':
                     v1 -= 1
                     v2 -= 1
+
                 # Assign bones to be joint1 - joint2, the pairs are pre-determined and hardcoded
                 # There also happens to be 25 bones
                 fp_sp[:, :, :, v1, :] = data[:, :, :, v1, :] - data[:, :, :, v2, :]

data_gen/ntu_gen_joint_data.py

@@ -3,11 +3,12 @@
 import pickle
 import argparse
 import numpy as np
+import yaml
 
 from tqdm import tqdm
 import sys
 sys.path.extend(['../'])
-from data_gen.preprocess import pre_normalization
+from preprocess import pre_normalization
 
 # For Cross-Subject benchmark "xsub"
 training_subjects = [1, 2, 4, 5, 8, 9, 13, 14, 15, 16, 17, 18, 19, 25, 27, 28, 31, 34, 35, 38]
@@ -72,8 +73,8 @@ def get_nonzero_std(s):
         s = 0
     return s
 
-
-def read_xyz(file, max_body=4, num_joint=25):  # 取了前两个body
+# 관절의 x,y,z 읽어오기
+def read_xyz(file, max_body=4, num_joint=25): 
     seq_info = read_skeleton_filter(file)
     # Create data tensor of shape: (# persons (M), # frames (T), # nodes (V), # channels (C))
     data = np.zeros((max_body, seq_info['numFrame'], num_joint, 3))
@@ -87,80 +88,156 @@ def read_xyz(file, max_body=4, num_joint=25):  # 取了前两个body
     energy = np.array([get_nonzero_std(x) for x in data])
     index = energy.argsort()[::-1][0:max_body_true]
     data = data[index]
-    # Data new shape: (C, T, V, M)
+
+    # Data new shape: (C, T, V, M) (# channels (C), # frames (T), # nodes (V), # persons (M))
     data = data.transpose(3, 1, 2, 0)
     return data
 
 
-def gendata(data_path, out_path, ignored_sample_path=None, benchmark='xview', part='eval'):
+def gendata(data_path, out_path, node_mask=None, attention_action=None, ignored_sample_path=None, 
+            benchmark='xview', part='eval'):
+
+    ########################################### 데이터 선별 ###########################################
+    # attention action, ignored_sample_path, benchmark, part에 따라 Train 데이터와 Test 데이터를 선별 
+    # 선별된 데이터의 파일 이름과 Label을 List로 저장
+    ###################################################################################################
+    # ignored_sample_path
     if ignored_sample_path != None:
         with open(ignored_sample_path, 'r') as f:
             ignored_samples = [line.strip() + '.skeleton' for line in f.readlines()]
     else:
         ignored_samples = []
 
-    sample_name = []
-    sample_label = []
+    sample_name = [] # benchmark와 part에 따라 선택되는 데이터 파일의 이름을 저장하는 List
+    sample_label = [] # benchmark와 part에 따라 선택되는 데이터 파일의 라벨을 저장하는 List
+
+    # Label data 생성
     for filename in os.listdir(data_path):
+        # ignored_sample_path가 입력으로 들어오면 그 데이터는 Pass
         if filename in ignored_samples:
             continue
+
+        # attention_action에 입력이 있으면 그 action class만 데이터 생성
+        if attention_action != None :
+            if filename[filename.find('A'):filename.find('A') + 4] not in attention_action :
+                continue
+
         action_class = int(filename[filename.find('A') + 1:filename.find('A') + 4])
         subject_id = int(filename[filename.find('P') + 1:filename.find('P') + 4])
         camera_id = int(filename[filename.find('C') + 1:filename.find('C') + 4])
 
+        # View와 Sub의 Train 데이터셋 분리하는 부분
         if benchmark == 'xview':
-            istraining = (camera_id in training_cameras)
+            istraining = (camera_id in training_cameras) # istraining은 bool type
         elif benchmark == 'xsub':
-            istraining = (subject_id in training_subjects)
+            istraining = (subject_id in training_subjects) # istraining은 bool type
         else:
             raise ValueError('Invalid benchmark provided: {}'.format(benchmark))
 
+        # Train과 Val에 따라 데이터 선택
         if part == 'train':
             issample = istraining
         elif part == 'val':
             issample = not (istraining)
         else:
             raise ValueError('Invalid data part provided: {}'.format(part))
 
-        if issample:
+        # 선택된 Sample들의 Filename과 라벨 데이터를 list로 저장
+        if issample: 
             sample_name.append(filename)
-            sample_label.append(action_class - 1)
-
+
+            # attention action이 있으면, 라벨들이 앞으로 당겨지게끔 설정
+            if attention_action != None :
+                sample_label.append(attention_action.index(filename[filename.find('A'):filename.find('A') + 4]))
+
+            # attention action이 없으면, filename 안의 Action number가 라벨
+            else:
+                sample_label.append(action_class - 1)
+
+    # 선택된 Sample들의 이름과 라벨 저장
     with open('{}/{}_label.pkl'.format(out_path, part), 'wb') as f:
         pickle.dump((sample_name, list(sample_label)), f)
 
-    # Create data tensor with shape (# examples (N), C, T, V, M)
+    ########################################## 데이터 전처리 ##########################################
+    # 위에서 선별된 File name에 맞는 데이터를 Load하여 N, C, T, V, M의 joint 데이터를 Load
+    # Joint data Normalize
+    # node_mask로 선택한 node만 npy 파일로 데이터 저장
+    ###################################################################################################
+    # Joint data 생성    
+    # Joint data를 저장할 임시 변수 생성. tensor with shape (# examples (N), C=3, T=300, V=25, M=2)
     fp = np.zeros((len(sample_label), 3, max_frame, num_joint, max_body_true), dtype=np.float32)
 
-    # Fill in the data tensor `fp` one training example a time
+    # Sample File에서 Joint Data를 읽어와서 임시 변수(fp)에 저장
     for i, s in enumerate(tqdm(sample_name)):
-        data = read_xyz(os.path.join(data_path, s), max_body=max_body_kinect, num_joint=num_joint)
-        fp[i, :, :data.shape[1], :, :] = data
-
-    fp = pre_normalization(fp)
+        # Data new shape: (C, T, V, M)
+        data = read_xyz(os.path.join(data_path, s), max_body=max_body_kinect, num_joint=num_joint) # (C, T, V, M)
+        fp[i, :, :data.shape[1], :, :] = data # (N, C, T, V, M)
+
+    # fp normalize
+    fp = pre_normalization(fp) # (N, C=3, T=300, V=25, M=2)
+
+
+    # 특정 노드만 가져오기
+    if node_mask != None : 
+        node_mask = node_mask - np.ones([len(node_mask)], dtype=int)
+        fp = fp[:, :, :, node_mask, :]
+
+    # Joint data 저장
     np.save('{}/{}_data_joint.npy'.format(out_path, part), fp)
 
 
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='NTU-RGB-D Data Converter.')
-    parser.add_argument('--data_path', default='../data/nturgbd_raw/nturgb+d_skeletons/')
+    parser.add_argument('--data_path', default='/mnt/disk2/data/private_data/NTU_RGB+D/nturgb+d_skeletons/nturgbd_skeletons_s001_to_s017/')
     parser.add_argument('--ignored_sample_path',
-                        default='../data/nturgbd_raw/samples_with_missing_skeletons.txt')
-    parser.add_argument('--out_folder', default='../data/ntu/')
+                        default='./data/nturgbd_raw/samples_with_missing_skeletons.txt')
+    parser.add_argument('--out_folder', default='./data/ntu/')
+    parser.add_argument(
+            '--node_mask',
+            default=None,
+            help='사용할 노드') # 사용할 노드 
+    parser.add_argument(
+            '--selected_action',
+            default=None,
+            help='action to train') # 사용할 Action 번호
+    parser.add_argument(
+            '--config',
+            default='./config/node_config.yaml',
+            help='path to the configuration file') # config 파일이 위치하는 디렉토리
 
     benchmarks = ['xsub', 'xview']
     parts = ['train', 'val']
     arg = parser.parse_args()
 
+    if arg.config is not None: 
+        with open(arg.config, 'r') as f:
+            default_arg = yaml.load(f) # config 파일에 들어있는 keys. config key로 명명
+
+        # 예외 처리 
+        key = vars(arg).keys() # parser에 들어있는 key값. default key로 명명
+        for k in default_arg.keys():
+            if k not in key: # config key가 default key에 들어있지 않으면
+                print('WRONG ARG: {}'.format(k)) # default key에 해당 키가 없음을 알림.
+                assert (k in key)
+
+        parser.set_defaults(**default_arg) # 임의의 개수의 Keyword arguments를 받아서 default key -> config key로 변경
+
+    arg = parser.parse_args() # config key가 반영된 argument
+
+    print("arg : ", arg)
+
     for b in benchmarks:
         for p in parts:
             out_path = os.path.join(arg.out_folder, b)
             if not os.path.exists(out_path):
                 os.makedirs(out_path)
             print(b, p)
             gendata(
-                arg.data_path,
-                out_path,
-                arg.ignored_sample_path,
-                benchmark=b,
-                part=p)
+                arg.data_path, # 전처리할 Skeleton data가 위치한 경로. NTU RGB 데이터를 의미
+                out_path, # 전처리된 데이터가 위치할 공통 경로.
+                node_mask=arg.node_mask, # node mask. 25개의 Joint를 다 사용하지 않을 경우 사용할 Joint를 List 형태로 넣으면 된다. 
+                attention_action = arg.selected_action, # 특정 Action에 대해서만 데이터 생성. 원하는 Action을 List 형태로 넣으면 된다.
+                ignored_sample_path = arg.ignored_sample_path, # 무시할 Sample 데이터가 적혀진 txt파일의 경로. 
+                benchmark=b, # xsub / xview
+                part=p) # train할지, validation할지 

data_gen/ntu_gen_motion_data.py

@@ -8,30 +8,35 @@
 datasets = {'ntu/xview', 'ntu/xsub'}
 parts = {'joint', 'bone'}
 
-
 def gen_motion_data():
-    for dataset in datasets:
-        for set in sets:
-            for part in parts:
-                fn = '../data/{}/{}_data_{}.npy'.format(dataset, set, part)
+    for dataset in datasets: # 'ntu/xview', 'ntu/xsub'
+        for set in sets: # 'train', 'val'
+            for part in parts: # 'joint', 'bone'
+                fn = './data/{}/{}_data_{}.npy'.format(dataset, set, part)
+
+                # Joint, Bone 파일 존재하지 않으면 for문 continue 실행
                 if not os.path.exists(fn):
                     print('Joint/bone data does not exist for {} {} set'.format(dataset, set))
                     continue
 
+                # Joint, Bone 파일이 존재하면 
                 print('Generating motion data for', dataset, set, part)
+
+                # Joint, Bone 데이터 Load
                 data = np.load(fn)
                 (N, C, T, V, M) = data.shape
                 fp_sp = open_memmap(
-                    '../data/{}/{}_data_{}_motion.npy'.format(dataset, set, part),
+                    './data/{}/{}_data_{}_motion.npy'.format(dataset, set, part),
                     dtype='float32',
                     mode='w+',
                     shape=data.shape)
 
+                # Motion데이터는 dt동안 Joint가 움직인 거리. m = v(t+1) - v(t) 
                 # Loop through frames and insert motion difference
                 for t in tqdm(range(T - 1)):
                     fp_sp[:, :, t, :, :] = data[:, :, t + 1, :, :] - data[:, :, t, :, :]
 
-                # Pad last frame with 0
+                # 마지막 한 프레임은 0으로 채움
                 fp_sp[:, :, T - 1, :, :] = 0