Safemotion Lib
Loading...
Searching...
No Matches
generate_pkl_data.py
Go to the documentation of this file.
1import pickle
2import numpy as np
3import copy
4import os
5
7from smdataset.AnnotationDataManager import AnnotationDataManager
8
9from smutils.utils_os import search_file, create_directory
10from smutils.utils_data import save_pkl_data
11
12
13def make_dummy_list(anno, num):
14 """
15 어노테이션 데이터를 복제해서 더미데이터 리스트를 생성하는 기능
16 args:
17 anno (dict): 복재하기 위한 어노테이션 데이터
18 num (int): 복제 수량
19 return (list[dict]): 더미데이터 리스트
20 """
21
22 #복사
23 tmp_anno = copy.deepcopy(anno)
24
25 #복사한 데이터의 키포인트 스코어를 0으로 셋팅함
26 for i in range(len(tmp_anno['keypoints'])):
27 tmp_anno['keypoints'][i][-1] = 0
28
29 tmp_anno['action_id']['action'] = -1 #더미데이터의 행동 설정
30 tmp_anno['valid'] = 0 #유효성 셋팅
31
32 dummy_list = []
33 for i in range(num):
34 dummy_list.append(copy.deepcopy(tmp_anno))
35 return dummy_list
36
37def check_invaild(anno_list, ratio):
38 """
39 데이터의 유효성 확인 기능
40 args:
41 anno_list (list[dict]): 확인하려는 데이터
42 valid : 유효성 파라미터, 1 or 0, 1이면 유효한 데이터, 0이면 유효하지 않는 데이터
43 ratio (float): 유효하지 않는 데이터 비율, 해당 비율보다 작으면 유효성 확인을 통과하지 못함
44 return (bool): 유효성 확인 결과, 유효성 확인을 통과하지 못하면 True, 통과하면 False
45 """
46 sum_valid = 0
47 for anno in anno_list:
48 sum_valid += anno['valid']
49
50 return sum_valid < len(anno_list)*ratio
51
52def calc_label(anno_list, ratio=0.3, valid_ratio=0.5):
53 """
54 데이터의 클래스를 결정하는 기능
55 args:
56 anno_list (list[dict]): 데이터
57 필수 데이터 : action_id에 action 키가 존재 해야함, 해당 키에 해당 프레임에서의 클래스 정보가 입력되어 있어야함
58 ratio (float): 라벨 결정을 위한 데이터의 시작 위치(%), 클래스를 결정하는데 사용하는 전체 범위는 데이터 중앙의 (1-2*ratio) *100 %가 됨
59 valid_ratio (float): 유효하지 않는 데이터 비율, 해당 비율보다 작으면 유효성 확인을 통과하지 못함, 클래스 결정을 위한 구간에서 유효성 검사
60 return (int): 데이터의 클래스, 유효성 검사나 클래스 확인이 어려운 경우 None을 반환함
61 """
62 l = len(anno_list) #데이터 전체 길이
63 start_idx = int(l*ratio) #클래스를 결정할 구간의 시작 위치
64 end_idx = l-start_idx #클래스를 결정할 구간의 끝 위치
65
66 #설정구간에 대한 클래스 리스트를 생성함, 빈도수 확인을 위함
67 label_list = []
68 for i in range(start_idx, end_idx):
69 anno = anno_list[i]
70
71 #유효성 확인, 0보다 작으면 해당 정보는 사용하지 않는다는 의미임
72 if anno['action_id']['action'] < 0:
73 continue
74
75 #유효성 확인, valid가 0이면 복사한 데이터라는 의미임
76 if anno['valid'] == 0:
77 continue
78
79 label_list.append(anno['action_id']['action'])
80
81 #유효성 확인
82 if len(label_list) < (end_idx - start_idx)*valid_ratio:
83 return None
84
85 #빈도수 확인
86 values, counts = np.unique(np.array(label_list), return_counts=True)
87 max_count = np.max(counts) #최대 빈도수
88 labels = values[counts == max_count] #최대 빈도수 클래스
89 labels = labels.tolist()
90
91 #최대 빈도수에 해당하는 클래스가 2개 이상이면 해당 샘플은 사용하지 않음
92 if len(labels) > 1:
93 return None
94
95 return labels[0]
96
97def gen_pkl_data(anno_list, h, w, image_folder):
98 """
99 학습을 위한 데이터 1개를 생성하는 기능
100 데이터의 최근 3개의 프레임으로 클래스를 결정함
101 args:
102 anno_list (list[dict]): 데이터
103 h (int): 원본 영상의 높이
104 w (int): 원본 영상의 너비
105 image_folder (str): 영상 폴더
106 return (dict): 학습데이터
107 keypoint (np.array): 스켈레톤 좌표
108 keypoint_score (np.array): 스켈레톤 스코어
109 frame_dir (str): 영상 폴더
110 img_shape (tuple): 원본 영상의 크기
111 original_shape (tuple): 원본 영상의 크기
112 total_frames (int): 데이터의 프레임 수
113 label (int): 샘플의 클래스 번호
114 """
115
116 #가장 마지막 데이터 확인, 가장 최근 데이터가
117 label = anno_list[-1]['action_id']['action']
118
119 #가장 마지막 데이터 유효성 확인
120 if label == -1:
121 return None
122
123 #마지막(최근) 3개 프레임이 동일한지 체크
124 len_data = len(anno_list)
125 for i in range(1, 3):
126 if anno_list[len_data-i-1]['action_id']['action'] != label:
127 return None
128
129 #유효성 검사
130 if check_invaild(anno_list, 0.5):
131 return None
132
133 #데이터 생성
134 pkl_data = dict()
135
136 kps = []
137 for anno in anno_list:
138 kps.append(anno['keypoints'])
139
140 kps = np.array([kps])
141 pkl_data['keypoint'] = kps[..., :2]
142 pkl_data['keypoint_score'] = kps[..., 2]
143 pkl_data['frame_dir'] = image_folder
144 pkl_data['img_shape'] = (h, w)
145 pkl_data['original_shape'] = (h, w)
146 pkl_data['total_frames'] = kps.shape[1]
147 pkl_data['label'] = anno_list[-1]['action_id']['action']
148
149 return pkl_data
150
151def gen_pkl_data_ver2(anno_list, h, w, image_folder):
152 """
153 학습을 위한 데이터 1개를 생성하는 기능
154 데이터의 중심 40% 구간에서 클래스를 결정함
155 args:
156 anno_list (list[dict]): 데이터
157 h (int): 원본 영상의 높이
158 w (int): 원본 영상의 너비
159 image_folder (str): 영상 폴더
160 return (dict): 학습데이터
161 keypoint (np.array): 스켈레톤 좌표
162 keypoint_score (np.array): 스켈레톤 스코어
163 frame_dir (str): 영상 폴더
164 img_shape (tuple): 원본 영상의 크기
165 original_shape (tuple): 원본 영상의 크기
166 total_frames (int): 데이터의 프레임 수
167 label (int): 샘플의 클래스 번호
168 """
169
170 #유효성 체크
171 if check_invaild(anno_list, 0.5):
172 return None
173
174 #클래스 결정, 데이터의 중심 구간에서 클래스를 결정함
175 label = calc_label(anno_list, 0.3)
176
177 #결정된 클래스의 유효성 체크
178 if label is None:
179 return None
180
181 #데이터 생성
182 pkl_data = dict()
183
184 kps = []
185 for anno in anno_list:
186 kps.append(anno['keypoints'])
187
188 kps = np.array([kps])
189 pkl_data['keypoint'] = kps[..., :2]
190 pkl_data['keypoint_score'] = kps[..., 2]
191 pkl_data['frame_dir'] = image_folder
192 pkl_data['img_shape'] = (h, w)
193 pkl_data['original_shape'] = (h, w)
194 pkl_data['total_frames'] = kps.shape[1]
195 pkl_data['label'] = label
196
197 return pkl_data
198
199def gen_pkl_data_ver3(anno_list, h, w, image_folder):
200 """
201 학습을 위한 데이터 1개를 생성하는 기능
202 행동 클래스와, 포즈 클래스를 설정함
203 행동 클래스는 데이터의 중심 40% 구간에서 클래스를 결정함
204 포즈 클래스는 데이터의 마지막 프레임으로 결정함
205 args:
206 anno_list (list[dict]): 데이터
207 h (int): 원본 영상의 높이
208 w (int): 원본 영상의 너비
209 image_folder (str): 영상 폴더
210 return (dict): 학습데이터
211 keypoint (np.array): 스켈레톤 좌표
212 keypoint_score (np.array): 스켈레톤 스코어
213 frame_dir (str): 영상 폴더
214 img_shape (tuple): 원본 영상의 크기
215 original_shape (tuple): 원본 영상의 크기
216 total_frames (int): 데이터의 프레임 수
217 label_action (int): 샘플의 행동 클래스 번호
218 label_pose (int): 샘플의 포즈 클래스 번호
219 """
220
221 #유효성 체크
222 if check_invaild(anno_list, 0.5):
223 return None
224
225 #마지막(최근) 3프레임에 대한 유효성 체크
226 len_data = len(anno_list)
227 for i in range(0, 3):
228 if anno_list[len_data-i-1]['valid'] == 0:
229 return None
230
231 #행동 라벨 결정
232 action_label = calc_label(anno_list, 0.3)
233
234 #행동 라벨 유효성 체크
235 if action_label is None:
236 return None
237
238 #포즈 라벨 결정
239 pose_label = anno_list[-1]['action_id']['pose']
240
241 #데이터 생성
242 pkl_data = dict()
243
244 kps = []
245 for anno in anno_list:
246 kps.append(anno['keypoints'])
247
248 kps = np.array([kps])
249 pkl_data['keypoint'] = kps[..., :2]
250 pkl_data['keypoint_score'] = kps[..., 2]
251 pkl_data['frame_dir'] = image_folder
252 pkl_data['img_shape'] = (h, w)
253 pkl_data['original_shape'] = (h, w)
254 pkl_data['total_frames'] = kps.shape[1]
255 pkl_data['label_action'] = action_label
256 pkl_data['label_pose'] = pose_label
257
258 return pkl_data
259
260def gen_pkl_data_dual(anno_list, h, w, image_folder, data_manager):
261
262 label = anno_list[-1]['action_id']['action']
263 if label == -1:
264 return None
265
266 len_data = len(anno_list)
267 for i in range(1, 3):
268 if anno_list[len_data-i-1]['action_id']['action'] != label:
269 return None
270
271 if check_invaild(anno_list, 0.5):
272 return None
273
274 pkl_data = dict()
275
276 kps = []
277 for anno in anno_list:
278 kps.append(anno['keypoints'])
279
280 bboxes = []
281 for anno in anno_list:
282 bboxes.append(anno['bbox'])
283
284 filename_list = []
285 for anno in anno_list:
286 file_name = data_manager.get_image_name(anno['image_id'])
287 filename_list.append(file_name)
288
289 kps = np.array([kps])
290 bboxes = np.array([bboxes])
291 pkl_data['keypoint'] = kps[..., :2]
292 pkl_data['keypoint_score'] = kps[..., 2]
293 pkl_data['frame_dir'] = image_folder
294 pkl_data['img_shape'] = (h, w)
295 pkl_data['original_shape'] = (h, w)
296 pkl_data['total_frames'] = kps.shape[1]
297 pkl_data['label'] = anno_list[-1]['action_id']['action']
298 pkl_data['start_index']=0
299 pkl_data['modality']='RGB'
300 pkl_data['keypoint_img_shape'] = (h, w)
301 pkl_data['filename_tmpl'] = filename_list
302 pkl_data['bbox'] = bboxes
303
304 return pkl_data
305
306def make_action_pkl_file(data_manager, k, mode=None):
307 """
308 하나의 클립에서 학습 데이터를 생성하는 기능
309 args:
310 data_manager: AnnotationDataManager 객체, 어노테이션 데이터파일로 초기화까지 진행
311 k (int): 하나의 데이터를 만드는데 사용하는 프레임 숫자
312 mode (str): 단일 라벨 데이터인지 다중 라벨 데이터인지 설정하는 파라미터
313 'sc' : 단일 라벨 데이터
314 'mc' : 다중 라벨 데이터
315 return (dict): 하나의 클립에 대한 학습 데이터 리스트, 학습 데이터는 추적 아이디별로 저장됨, 추적 아이디의 학습 데이터는 list[dict]의 구조임
316 """
317
318 pkl_datas = {}
319
320 # h, w = data_manager.read_image().shape[:2]
321 h = 1080
322 w = 1920
323 image_folder = data_manager.get_image_folder()
324
325 #시작 프레임, 종료 프레임 설정
326 frames = data_manager.get_number_of_image()
327 start_frame = 0
328 end_frame = frames
329
330 #추적 아이디를 기준으로 데이터 형태 변환
331 track_data = data_manager.get_track_anno_data(filter=True, start_frame=start_frame, end_frame=end_frame)
332
333 #데이터 생성, 추적 아이디마다 생성함
334 for track_id, anno_list in track_data.items():
335
336 frame_no = len(anno_list)
337
338 image_idx_1st = data_manager.image_ids.index(anno_list[0]['image_id']) #시작 프레임
339 image_idx_end = data_manager.image_ids.index(anno_list[-1]['image_id']) #끝 프레임
340 interval = image_idx_end - image_idx_1st #데이터 길이
341
342 #유효성 체크, 데이터의 길이가 설정한 학습데이터의 길이보다 작을 경우 사용하지 않음
343 if interval < k:
344 continue
345
346 #라벨 변환
347 if mode == 'sc':
348 simple_anno_list = cvt_labelmap_v21(anno_list, cvt_labelmap_v21_action_to_simple_ver, cvt_labelmap_v21_pose_to_simple_ver)
349 elif mode == 'mc':
350 simple_anno_list = cvt_labelmap_v21_pose_and_action(anno_list, cvt_labelmap_v21_action_to_simple_ver2, cvt_labelmap_v21_pose_to_simple_ver2)
351
352 #유효성 셋팅
353 for i in range(0, frame_no):
354 simple_anno_list[i]['valid'] = 1
355
356 #추적 아이디의 데이터를 저장할 리스트 생성
357 pkl_datas[track_id] = []
358
359 #pkl data 생성
360 anno_q = [simple_anno_list[0]]
361 for i in range(1, frame_no):
362 #인접한 데이터의 프레임 확인
363 image_idx_end = data_manager.image_ids.index(simple_anno_list[i]['image_id'])
364 image_idx_end_m1 = data_manager.image_ids.index(anno_q[-1]['image_id'])
365
366 #인접한 데이터의 프레임 간격에 따라서 데이터 처리
367 interval = image_idx_end - image_idx_end_m1
368 if interval > k: #k보다 클경우 다른 클립으로 분리함
369 anno_q = []
370 elif interval > 1: #k보다 작고 1보다 클경우 더미 데이터를 넣어줌
371 anno_q.extend(make_dummy_list(anno_q[-1], interval-1))
372
373 #큐에 데이터 붙여줌
374 anno_q.append(simple_anno_list[i])
375
376 #큐의 데이터가 k개 이상인 경우만 학습 데이터 생성 프로세스를 진행함
377 frame_num = len(anno_q)
378 if frame_num < k:
379 continue
380
381 #모드에 따라 학습 데이터 생성
382 if mode == 'sc':
383 pkl_data = gen_pkl_data_ver2(anno_q[frame_num-k:], h, w, image_folder)
384 elif mode == 'mc':
385 pkl_data = gen_pkl_data_ver3(anno_q[frame_num-k:], h, w, image_folder)
386 #pkl_data = gen_pkl_data(simple_anno_list[i-(k-1):i+1], h, w, image_folder)
387
388 #데이터가 생성되지 않았을 경우
389 if pkl_data is None:
390 continue
391
392 #추적 아이디의 키에 피클 데이터 저장
393 pkl_datas[track_id].append(pkl_data)
394
395 #추적 아이디의 학습 데이터가 생성되지 않을 경우 키 제거
396 if len(pkl_datas[track_id]) == 0:
397 del pkl_datas[track_id]
398
399 return pkl_datas
400
401def make_action_pkl_file_dual(data_manager, k):
402 # k : 하나의 데이터를 만드는데 드는 프레임 숫자
403 pkl_datas = []
404
405 h, w = data_manager.read_image().shape[:2]
406 image_folder = data_manager.get_image_folder()
407
408 frames = data_manager.get_number_of_image()
409
410 track_data = data_manager.get_track_anno_data(filter=True, start_frame=start_frame, end_frame=end_frame)
411 for track_id, anno_list in track_data.items():
412 frame_no = len(anno_list)
413
414 image_idx_1st = data_manager.image_ids.index(anno_list[0]['image_id'])
415 image_idx_end = data_manager.image_ids.index(anno_list[-1]['image_id'])
416 interval = image_idx_end - image_idx_1st
417 if interval < k:
418 continue
419
420 simple_anno_list = cvt_labelmap_v21(anno_list, cvt_labelmap_v21_action_to_simple_ver, cvt_labelmap_v21_pose_to_simple_ver)
421
422 for i in range(0, frame_no):
423 simple_anno_list[i]['valid'] = 1
424
425 anno_q = [simple_anno_list[0]]
426
427 #file_list = os.listdir(image_folder)
428 #file_list = sorted(file_list)
429
430 for i in range(1, frame_no):
431 image_idx_end = data_manager.image_ids.index(simple_anno_list[i]['image_id'])
432 image_idx_end_m1 = data_manager.image_ids.index(anno_q[-1]['image_id'])
433
434 interval = image_idx_end - image_idx_end_m1
435 if interval > 30:
436 anno_q = []
437 elif interval > 1:
438 anno_q.extend(make_dummy_list(anno_q[-1], interval-1))
439 anno_q.append(simple_anno_list[i])
440
441 frame_num = len(anno_q)
442 if frame_num < k:
443 continue
444
445 pkl_data = gen_pkl_data_dual(anno_q[frame_num-k:], h, w, image_folder, data_manager)
446 #pkl_data = gen_pkl_data(simple_anno_list[i-(k-1):i+1], h, w, image_folder)
447 if pkl_data is None:
448 continue
449
450 pkl_datas.append(pkl_data)
451
452 return pkl_datas
453
454def split_from_action(data_list):
455 """
456 학습 데이터를 행동 클래스에 따라서 분리하는 기능
457 args:
458 data_list (list[dict]): 학습 데이터 리스트, 데이터에는 label or label_action 키에 대한 값이 존재해야함
459 return (list[dict]) : 행동 클래스를 기준으로 분리된 학습 데이터
460 """
461 split_data = {}
462 for data in data_list:
463 if 'label' in data:
464 label = data['label']
465 else:
466 label = data['label_action']
467
468 if label not in split_data:
469 split_data[label] = []
470 split_data[label].append(data)
471 return split_data
472
473def split_from_pose(data_list):
474 """
475 학습 데이터를 포즈 클래스에 따라서 분리하는 기능
476 args:
477 data_list (list[dict]): 학습 데이터 리스트, 데이터에는 label_pose 키에 대한 값이 존재해야함
478 return (list[dict]) : 포즈 클래스를 기준으로 분리된 학습 데이터
479 """
480 split_data = {}
481 for data in data_list:
482 label = data['label_pose']
483
484 if label not in split_data:
485 split_data[label] = []
486 split_data[label].append(data)
487 return split_data
488
489
491 """
492 클립 정보를 담고있는 어노테이션 파일들을 사용해서 학습데이터를 만드는 기능
493 args:
494 pkl_root (str): 학습데이터를 저장할 폴더
495 dataset_folder (str): 클립 정보를 담고있는 어노테이션 파일들이 저장된 폴더
496 k (int): 학습데이터를 생성하는 프레임 수
497 mode (str): 단일 라벨 데이터인지 다중 라벨 데이터인지 설정하는 파라미터
498 'sc' : 단일 라벨 데이터
499 'mc' : 다중 라벨 데이터
500 """
501 data_manager = AnnotationDataManager()
502
503 create_directory(args.pkl_root)
504
505 json_name_list, json_path_list = search_file(args.dataset_folder, '.json')
506 file_num = len(json_name_list)
507
508 for i, (json_path, json_name) in enumerate(zip(json_path_list, json_name_list)):
509 data_manager.load_annotation(json_path)
510 data_manager.init_annotation()
511
512 pkl_datas = make_action_pkl_file(data_manager, args.k, args.mode)
513 #pkl_datas = make_action_pkl_file_dual(data_manager, args.k)
514
515 for track_id, datas in pkl_datas.items():
516 split_data = split_from_action(datas)
517
518 for action_label, data in split_data.items():
519 pkl_name = json_name.replace('json', 'pkl') #name 변경
520 name, ext = pkl_name.split('.') #name 분리 name = P001C021A0302R001, ext = pkl
521 if args.mode == 'mc':
522 s_data = split_from_pose(data)
523 for pose_label, d in s_data.items():
524 new_pkl_name = f'{name}_{track_id:02d}_{action_label:03d}_{pose_label:03d}.{ext}' #pkl_name 재설정
525 save_pkl_path = os.path.join(args.pkl_root, new_pkl_name) # 저장경로 설정
526 save_pkl_data(d, save_pkl_path) #저장
527
528 #저장 데이터 출력
529 print(f"[mc][ {i+1} / {file_num} ] {save_pkl_path} : {len(d)}")
530
531 elif args.mode == 'sc':
532 new_pkl_name = f'{name}_{track_id:02d}_{action_label:03d}.{ext}' #pkl_name 재설정
533 save_pkl_path = os.path.join(args.pkl_root, new_pkl_name) # 저장경로 설정
534 save_pkl_data(data, save_pkl_path) #저장
535
536 #저장 데이터 출력
537 print(f"[sc][ {i+1} / {file_num} ] {save_pkl_path} : {len(datas)}")
538 else:
539 print('error')
540
541
542def merge_pkl(folder, file_path_list=None ):
543 """
544 학습데이터 파일들을 하나로 합치는 기능
545 args:
546 folder (str): 학습데이터가 저장된 폴더
547 file_path_list (list[str]) : 학습데이터 경로 리스트, 설정하지 않으면 folder에서 찾음
548 return (list[dict]): 합쳐진 학습 데이터
549 """
550 if file_path_list == None:
551 file_name_list, file_path_list = search_file(folder, fileEx='.pkl')
552
553 merged_pkl_datas = []
554 for path in file_path_list:
555 with open(path, 'rb') as f:
556 pkl_data = pickle.load(f)
557 merged_pkl_datas += pkl_data
558 print(f"{path} : {len(pkl_data)}")
559
560 return merged_pkl_datas
561
563 """
564 학습데이터 파일들을 하나로 합치고 저장하는 기능
565 args:
566 pkl_root (str): 학습데이터가 저장된 폴더
567 merge_pkl_path (str): 합쳐진 데이터를 하나의 파일로 저장하는 경로
568 """
569 merged_pkl_datas = merge_pkl(args.pkl_root)
570 print(f"total {args.merge_pkl_path} : {len(merged_pkl_datas)}")
571 save_pkl_data(merged_pkl_datas, args.merge_pkl_path)
572
573
574def check_label_v22(anno_list, category, target_label, ratio=0.3):
575 """
576 데이터에 특정 카테고리의 특정 클래스가 존재하는지 확인하는 기능
577 args:
578 anno_list (list[dict]): 데이터
579 필수 데이터 : action_id에 라벨 이름(category)에 대한 키가 존재 해야함, 해당 키에 해당 프레임에서의 클래스 정보가 입력되어 있어야함
580 category (str): 데이터에서 확인하려는 카테고리
581 target_label (int): 데이터에서 확인하려는 클래스
582 ratio (float): 라벨 확인을 위한 데이터의 시작 위치(%), 클래스를 확인하는데 사용하는 전체 범위는 데이터 중앙의 (1-2*ratio) *100 %가 됨
583
584 return (bool): 원하는 클래스 존재 유무
585 """
586
587 #확인 구간 설정
588 l = len(anno_list)
589 start_idx = int(l*ratio)
590 end_idx = l-start_idx
591
592 #확인 구간의 클래스 라벨 리스트 생성, 빈도수를 확인하고
593 label_list = []
594 for i in range(start_idx, end_idx):
595 anno = anno_list[i]
596
597 #유효성 체크, 0 이하로 클래스가 저장되어 있으면 해당 프레임은 사용하지 않는 다는 의미임
598 if anno['action_id'][category] < 0:
599 continue
600
601 #유효성 체크, valid가 0이면 더미데이터를 의미함
602 if anno['valid'] == 0:
603 continue
604
605 label_list.append(anno['action_id'][category])
606
607 return target_label in label_list
608
609 #빈도수 기반으로 라벨 체크
610 # values, counts = np.unique(np.array(label_list), return_counts=True)
611 # if target_label not in list(values):
612 # return False
613 # return True
614
615def gen_pkl_data_v22(anno_list, category, target_label, h, w, image_folder):
616 """
617 학습을 위한 데이터 1개를 생성하는 기능
618 특정 카테고리의 특정 클래스에 해당하는 데이터를 생성함
619 데이터의 중심 40% 구간을 확인해서 클래스를 설정함
620 args:
621 anno_list (list[dict]): 데이터
622 category (str): 데이터에서 확인하려는 카테고리
623 target_label (int): 데이터에서 확인하려는 클래스
624 h (int): 원본 영상의 높이
625 w (int): 원본 영상의 너비
626 image_folder (str): 영상 폴더
627 return (dict): 학습데이터
628 keypoint (np.array): 스켈레톤 좌표
629 keypoint_score (np.array): 스켈레톤 스코어
630 frame_dir (str): 영상 폴더
631 img_shape (tuple): 원본 영상의 크기
632 original_shape (tuple): 원본 영상의 크기
633 total_frames (int): 데이터의 프레임 수
634 category (str): 샘플의 카테고리
635 label (int): 샘플의 클래스 번호
636 """
637
638 #유효성 체크
639 if check_invaild(anno_list, 0.5):
640 return None
641
642 #카테고리의 클래스 체크, 클래스가 중심 40% 구간에 존재하지 않으면 학습 데이터를 생성하지 않음
643 if check_label_v22(anno_list, category, target_label, 0.3):
644 label = target_label
645 else:
646 return None
647
648 #데이터 생성
649 pkl_data = dict()
650
651 kps = []
652 for anno in anno_list:
653 kps.append(anno['keypoints'])
654
655 kps = np.array([kps])
656 pkl_data['keypoint'] = kps[..., :2]
657 pkl_data['keypoint_score'] = kps[..., 2]
658 pkl_data['frame_dir'] = image_folder
659 pkl_data['img_shape'] = (h, w)
660 pkl_data['original_shape'] = (h, w)
661 pkl_data['total_frames'] = kps.shape[1]
662 pkl_data['category'] = category
663 pkl_data['label'] = label
664
665 return pkl_data
666
667def make_action_pkl_file_v22(data_manager, category, target_label, k):
668 """
669 하나의 클립에서 학습 데이터를 생성하는 기능
670 카테고리의 특정 클래스에 대한 데이터를 생성함
671 args:
672 data_manager: AnnotationDataManager 객체, 어노테이션 데이터파일로 초기화까지 진행
673 category (str): 데이터에서 확인하려는 카테고리
674 target_label (int): 데이터에서 확인하려는 클래스
675 k (int): 하나의 데이터를 만드는데 사용하는 프레임 숫자
676 return (dict): 하나의 클립에 대한 학습 데이터 리스트, 학습 데이터는 추적 아이디별로 저장됨, 추적 아이디의 학습 데이터는 list[dict]의 구조임
677 """
678
679 #이미지 크기 및 폴더 셋팅
680 # h, w = data_manager.read_image().shape[:2]
681 h = 1080
682 w = 1920
683 image_folder = data_manager.get_image_folder()
684
685 #클립의 시작, 끝 프레임 설정
686 start_frame = data_manager.dataset['info']['start']
687 end_frame = data_manager.dataset['info']['end']
688
689 #추적 아이디를 기준으로 데이터 형태 변환
690 track_data = data_manager.get_track_anno_data(filter=True, start_frame=start_frame, end_frame=end_frame)
691
692 #학습 데이터 생성, 추적 아이디별로 생성함
693 pkl_datas = {}
694 for track_id, anno_list in track_data.items():
695
696 frame_no = len(anno_list)
697
698 image_idx_1st = data_manager.image_ids.index(anno_list[0]['image_id']) #시작 프레임
699 image_idx_end = data_manager.image_ids.index(anno_list[-1]['image_id']) #끝 프레임
700
701 #유효성 체크, 데이터의 길이가 설정한 학습데이터의 길이보다 작을 경우 사용하지 않음
702 interval = image_idx_end - image_idx_1st
703 if interval < k:
704 continue
705
706 #학습 클래스 변환, 라벨링된 데이터를 학습을 위한 데이터로 변환
707 simple_anno_list = cvt_labelmap_v22(anno_list)
708
709 #유효성 셋팅
710 for i in range(0, frame_no):
711 simple_anno_list[i]['valid'] = 1
712
713 #추적 아이디의 데이터를 저장할 리스트 생성
714 pkl_datas[track_id] = []
715
716 #pkl data 생성
717 anno_q = [simple_anno_list[0]]
718 for i in range(1, frame_no):
719
720 #인접한 데이터의 프레임 확인
721 image_idx_end = data_manager.image_ids.index(simple_anno_list[i]['image_id'])
722 image_idx_end_m1 = data_manager.image_ids.index(anno_q[-1]['image_id'])
723
724 #인접한 데이터의 프레임 간격에 따라서 데이터 처리
725 interval = image_idx_end - image_idx_end_m1
726 if interval > k: #k보다 클경우 다른 클립으로 분리함
727 anno_q = []
728 elif interval > 1: #k보다 작고 1보다 클경우 더미 데이터를 넣어줌
729 anno_q.extend(make_dummy_list(anno_q[-1], interval-1))
730
731 #큐에 데이터를 붙여줌
732 anno_q.append(simple_anno_list[i])
733
734 #큐의 데이터가 k개 이상인 경우만 학습 데이터 생성 프로세스를 진행함
735 frame_num = len(anno_q)
736 if frame_num < k:
737 continue
738
739 #학습 데이터 생성
740 pkl_data = gen_pkl_data_v22(anno_q[frame_num-k:], category, target_label, h, w, image_folder)
741
742 #학습 데이터가 생성되지 않았을 경우
743 if pkl_data is None:
744 continue
745
746 #추적 아이디의 키에 피클 데이터 저장
747 pkl_datas[track_id].append(pkl_data)
748
749 #추적 아이디의 학습 데이터가 없을 경우 키 제거
750 if len(pkl_datas[track_id]) == 0:
751 del pkl_datas[track_id]
752
753 return pkl_datas
754
755
757 """
758 클립 정보를 담고있는 어노테이션 파일들을 사용해서 학습데이터를 만드는 기능
759 args:
760 category_k (dict): 카테고리별 학습데이터를 생성하는 프레임 수
761 data_root (str): 클립 정보를 담고있는 어노테이션 파일들이 저장된 폴더
762 pkl_root (str): 학습데이터를 저장할 폴더
763 """
764
765 #어노테이션 관리자 설정
766 data_manager = AnnotationDataManager()
767
768 #학습데이터 생성을 위한 카테고리별 프레임 수
769 category_k = args.category_k
770
771 data_root = args.data_root #클립 저장 폴더
772 for category, k in category_k.items():
773 #클립 리스트
774 dataset_folder = os.path.join(data_root, category)
775 json_name_list, json_path_list = search_file(dataset_folder, '.json')
776
777 #카테고리 별 저장 폴더 생성
778 save_folder = os.path.join(args.pkl_root, category)
779 for i, (json_path, json_name) in enumerate(zip(json_path_list, json_name_list)):
780 target_label = json_path.split('/')[-2]
781 cat_label = os.path.join(save_folder, target_label)
782 create_directory(cat_label)
783
784 file_num = len(json_name_list) #클립 수량, 콘솔 출력용
785 total_data = 0 #학습 데이터 수량, 콘솔 출력용
786 for i, (json_path, json_name) in enumerate(zip(json_path_list, json_name_list)):
787 # print(f"{category} -> [ {i+1} / {file_num} ] {json_name}")
788
789 #어노테이션 로드
790 data_manager.load_annotation(json_path)
791 data_manager.init_annotation()
792
793 #타겟 라벨 -> 어노테이션 하나에 하나의 클래스에 대해서만 학습데이터를 생성함
794 target_label = int(json_path.split('/')[-2])
795
796 #학습 데이터 생성
797 pkl_datas = make_action_pkl_file_v22(data_manager, category, target_label, k)
798
799 #학습 데이터 저장
800 for track_id, datas in pkl_datas.items():
801 pkl_name = json_name.replace('json', 'pkl')
802 save_pkl_path = os.path.join(save_folder, f'{target_label:02d}', pkl_name)
803 save_pkl_data(datas, save_pkl_path)
804 total_data += len(datas)
805 print(f"{category} -> [ {i+1} / {file_num} ] {save_pkl_path} : {len(datas):5d}/{total_data:}")
make_action_pkl_file_v22(data_manager, category, target_label, k)
gen_pkl_data_ver3(anno_list, h, w, image_folder)
gen_pkl_data_v22(anno_list, category, target_label, h, w, image_folder)
run_generate_pkl_data_from_json_v22(args)
calc_label(anno_list, ratio=0.3, valid_ratio=0.5)
check_label_v22(anno_list, category, target_label, ratio=0.3)
make_action_pkl_file(data_manager, k, mode=None)
gen_pkl_data_dual(anno_list, h, w, image_folder, data_manager)
gen_pkl_data_ver2(anno_list, h, w, image_folder)
check_invaild(anno_list, ratio)
gen_pkl_data(anno_list, h, w, image_folder)
make_action_pkl_file_dual(data_manager, k)
merge_pkl(folder, file_path_list=None)