-
Notifications
You must be signed in to change notification settings - Fork 2
/
tracking_data_assembler.py
74 lines (61 loc) · 2.49 KB
/
tracking_data_assembler.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
import cv2
import h5py
import numpy as np
import tfplus
import tfplus.utils.progress_bar as pb
class TrackingDataAssembler(object):
"""
{video_id}/video/frm_{frame_id}/image: PNG encoded image
{video_id}/annotations/obj_{object_id}/bbox: {left, top, right, bottom} * num_valid_frames
{video_id}/annotations/obj_{object_id}/presence: frame indices of appearance
"""
def __init__(self, output_fname):
self.log = tfplus.utils.logger.get()
self.output_fname = output_fname
self.log.info('Output h5 dataset: {}'.format(self.output_fname))
self.log.info('Reading video IDs')
self.vid_ids = self.get_video_ids()
self.log.info('Video IDs: {}'.format(self.vid_ids))
pass
def save(self, key, data, group):
if key in group:
del group[key]
group[key] = data
pass
def encode(self, img):
return cv2.imencode('.png', img)[1]
def get_video_ids(self):
raise Exception('Not implemented')
def get_frame_ids(self, vid_id):
raise Exception('Not implemented')
def get_frame_img(self, vid_id, frm_id):
raise Exception('Not implemented')
def get_obj_ids(self, vid_id):
raise Exception('Not implemented')
def get_obj_data(self, vid_id, obj_id):
raise Exception('Not implemented')
def assemble(self):
num_vid = len(self.vid_ids)
with h5py.File(self.output_fname, 'a') as h5f:
for idx in pb.get(num_vid):
vid_id = self.vid_ids[idx]
for frm_id in self.get_frame_ids(vid_id):
frm_img = self.get_frame_img(vid_id, frm_id)
img_enc = self.encode(frm_img)
frm_key = '{}/video/frm_{}/image'.format(vid_id, frm_id)
self.save(frm_key, img_enc, h5f)
obj_ids = self.get_obj_ids(vid_id)
if obj_ids is None:
# Testing data do not have labels.
continue
for obj_id in obj_ids:
obj_data = self.get_obj_data(vid_id, obj_id)
obj_key = '{}/annotations/obj_{}'.format(vid_id, obj_id)
frm_nonzero = obj_data['presence'].nonzero()[0]
self.save(obj_key + '/bbox',
obj_data['bbox'][frm_nonzero], h5f)
self.save(obj_key + '/frame_indices', frm_nonzero, h5f)
pass
pass
pass
pass