-
Notifications
You must be signed in to change notification settings - Fork 1
/
annotate_actions.py
109 lines (92 loc) · 4.04 KB
/
annotate_actions.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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
from glob import glob
import cv2
import numpy as np
import matplotlib.pyplot as plt
import scipy.io as sio
import os
import pickle
# helper function wrapping cv2.putText with default values
def puttext(im, title, coords, size=0.4, color=(255, 255, 255)):
cv2.putText(im, title, coords, cv2.FONT_HERSHEY_SIMPLEX, size, color)
return
def save_tube_annot(tube_annot):
if not os.path.isdir(out_folder + '/' + video_name):
os.makedirs(out_folder + '/' + video_name)
tube_save_path = out_folder + '/' + video_name + '/' + str(tube_id) + '.pickle'
# np.save(tube_save_path, tube_annot)
with open(tube_save_path, 'wb') as handle:
pickle.dump(tube_annot, handle, protocol=pickle.HIGHEST_PROTOCOL)
return
cls = 'PoleVault'
annot_types = {'1': 'pole_down', '2': 'over_bar', '3': 'end_jump'}
dataset_path = 'video/'
annot_path = 'annots/'
out_folder = 'new_annots/'
annots = sio.loadmat(annot_path + 'finalAnnots', struct_as_record=True)['annot'][0] # all annotations
cur_annots = [x for x in annots if cls == x[1][0].split('/')[0]] # take only annotations of current class
draw_bb = False
for a in cur_annots:
num_frames = a[0]
video_name = a[1][0]
tubes = a[2][0]
num_tubes = len(tubes)
for tube_id in range(num_tubes):
start_frame = tubes['sf'][tube_id][0][0]
end_frame = tubes['ef'][tube_id][0][0]
boxes = tubes['boxes'][tube_id].astype(np.int32)
assert(boxes.shape[0] == end_frame - start_frame + 1)
frame_id = start_frame
done = False
if os.path.isfile(out_folder + '/' + video_name + '/' + str(tube_id) + '.pickle'):
# tube_annot = np.load(out_folder + '/' + video_name + '/' + str(tube_id) + '.npy')
with open(out_folder + '/' + video_name + '/' + str(tube_id) + '.pickle', 'rb') as handle:
tube_annot = pickle.load(handle)
else:
tube_annot = {}
# frame names are numerated matlab style
while not done:
# frame index
index_in_tube = frame_id - start_frame
frame_key = video_name + '/' + str(tube_id) + '_' + str(index_in_tube) + '_' + str(frame_id)
cur_action = None
annot_items = tube_annot.items() # key, val
keys_ = [a[0] for a in annot_items]
vals_ = [a[1] for a in annot_items]
if [frame_id, index_in_tube] in vals_:
cur_action = (keys_[np.where([x == [frame_id, index_in_tube] for x in vals_])[0][0]])
# read image
im = cv2.imread(dataset_path + video_name + '/' + str(frame_id).zfill(5) + '.jpg')
# show bounding box
if draw_bb:
box = boxes[index_in_tube]
cv2.rectangle(im, (box[0], box[1]), (box[0]+box[2], box[1]+box[3]), (0,255,0), 3)
# shitty way to add a title to an image with opencv
im = np.pad(im, ((40,0), (0,0), (0,0)), 'constant', constant_values=0)
puttext(im, video_name + ' - tube: ' + str(tube_id), (int(im.shape[0]/2) - 100, 15))
puttext(im, 'FRAME ' + str(index_in_tube), (int(im.shape[0]/2) - 100, 35))
if cur_action is not None:
puttext(im, cur_action, (int(im.shape[0]/2), 35))
cv2.imshow('frame', im)
key = cv2.waitKey(0)
# next frame
if key == ord('d'):
frame_id = min(frame_id + 1, end_frame)
# prev frame
elif key == ord('a'):
frame_id = max(frame_id - 1, start_frame)
# draw boxes
elif key == ord('b'):
draw_bb = not draw_bb
# next tube
elif key == 13:
save_tube_annot(tube_annot)
done = True
# numbers 0-9
elif 48 <= key <= 57:
key = str(key - 48)
if key in annot_types.keys():
tube_annot[annot_types[key]] = [frame_id, index_in_tube]
# ESC
elif key == 27:
print('bye')
raise SystemExit(0)