-
Notifications
You must be signed in to change notification settings - Fork 90
/
Annotator.py
333 lines (265 loc) · 11 KB
/
Annotator.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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
# -*- coding: utf-8 -*-
# Copyright 2018-2019 Jacob M. Graving <jgraving@gmail.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
import h5py
import os
from deepposekit.annotate.gui.GUI import GUI
from deepposekit.annotate.utils import hotkeys as keys
__all__ = ["Annotator"]
class Annotator(GUI):
"""
A GUI for annotating images.
------------------------------------------------------------
Keys | Action
------------------------------------------------------------
> +,- | Rescale the image
> Left mouse | Move active keypoint
> W, A, S, D | Move active keypoint
> space | Changes W,A,S,D mode (swaps between 1px or 10px)
> J, L | Load previous or next image
> <, > | Jump 10 images backward or forward
> I, K or |
tab, shift+tab | Switch active keypoint
> R | Mark frame as unannotated, or "reset"
> F | Mark frame as annotated or "finished"
> Esc, Q | Quit the Annotator GUI
------------------------------------------------------------
Note: Data is automatically saved when moving between frames.
Parameters
----------
datapath: str
Filepath of the HDF5 (.h5) file that contains the images to
be annotated.
dataset: str
Key name to access the images in the .h5 file.
skeleton: str
Filepath of the .csv or .xlsx file that has indexed information
on name of the keypoint (part, e.g. head), parent (the direct
connecting part, e.g. neck connects to head, parent is head),
and swap (swapping positions with a part when reflected).
See example file for more information.
scale: int/float, default 1
Scaling factor for the GUI (e.g. used in zooming).
text_scale: float
Scaling factor for the GUI font.
A text_scale of 1 works well for 1920x1080 (1080p) images
shuffle_colors: bool, default = True
Whether to shuffle the color order for drawing keypoints
refresh: int, default 100
Delay on receiving next keyboard input in milliseconds.
Attributes
----------
window_name: str
Name of the Annotation window when running program.
Set to be 'Annotation' unless otherwise changed.
n_images: int
Number of images in the .h5 file.
n_keypoints: int
Number of keypoints in the skeleton.
key: int
The key that is pressed on the keyboard.
image_idx: int
Index of a specific image in the .h5 file.
image: numpy.ndarray
One image accessed using image_idx.
Example
-------
>>> from deepposekit import Annotator
>>> app = Annotator('annotation.h5', 'images', 'skeleton.csv')
>>> app.run()
"""
def __init__(
self,
datapath,
dataset,
skeleton,
scale=1,
text_scale=0.15,
shuffle_colors=True,
refresh=100,
):
super(GUI, self).__init__()
self.window_name = "Annotation"
self.shuffle_colors = shuffle_colors
self._init_skeleton(skeleton)
if os.path.exists(datapath):
self._init_data(datapath, dataset)
else:
raise ValueError("datapath file or path does not exist")
self._init_gui(scale, text_scale, shuffle_colors, refresh)
def _init_data(self, datapath, dataset):
""" Initializes the images from the .h5 file (called in init).
Parameters
----------
datapath: str
Path of the .h5 file that contains the images to be annotated.
dataset: str
Key name to access the images in the .h5 file.
"""
if isinstance(datapath, str):
if datapath.endswith(".h5"):
self.datapath = datapath
else:
raise ValueError("datapath must be .h5 file")
else:
raise TypeError("datapath must be type `str`")
if isinstance(dataset, str):
self.dataset = dataset
else:
raise TypeError("dataset must be type `str`")
with h5py.File(self.datapath, "r+") as h5file:
self.n_images = h5file[self.dataset].shape[0]
# Check that all parts of the file exist
if "annotations" not in list(h5file.keys()):
empty_array = np.zeros((self.n_images, self.n_keypoints, 2))
h5file.create_dataset(
"annotations",
(self.n_images, self.n_keypoints, 2),
dtype=np.float64,
data=empty_array,
)
for idx in range(self.n_images):
h5file["annotations"][idx] = self.skeleton.loc[:, ["x", "y"]].values
if "annotated" not in list(h5file.keys()):
empty_array = np.zeros((self.n_images, self.n_keypoints), dtype=bool)
h5file.create_dataset(
"annotated",
(self.n_images, self.n_keypoints),
dtype=bool,
data=empty_array,
)
if "skeleton" not in list(h5file.keys()):
skeleton = self.skeleton[["tree", "swap_index"]].values
h5file.create_dataset(
"skeleton", skeleton.shape, dtype=np.int32, data=skeleton
)
# Unpack the images from the file
self.image_idx = np.sum(np.all(h5file["annotated"].value, axis=1)) - 1
self.image = h5file[self.dataset][self.image_idx]
self._check_grayscale()
self.skeleton.loc[:, ["x", "y"]] = h5file["annotations"][self.image_idx]
self.skeleton.loc[:, "annotated"] = h5file["annotated"][self.image_idx]
def _save(self):
""" Saves an image.
Automatically called when moving to new images or invoked manually
using 'ctrl + s' keys.
"""
with h5py.File(self.datapath) as h5file:
h5file["annotations"][self.image_idx] = self.skeleton.loc[
:, ["x", "y"]
].values
h5file["annotated"][self.image_idx] = self.skeleton.loc[
:, "annotated"
].values
self.skeleton.loc[:, ["x", "y"]] = h5file["annotations"][self.image_idx]
self.skeleton.loc[:, "annotated"] = h5file["annotated"][self.image_idx]
def _load(self):
""" Loads an image.
This method is called in _move_image_idx when moving to different
images. The image of specified image_idx will be loaded onto the GUI.
"""
with h5py.File(self.datapath) as h5file:
self.image = h5file[self.dataset][self.image_idx]
self._check_grayscale()
self.skeleton.loc[:, ["x", "y"]] = h5file["annotations"][self.image_idx]
self.skeleton.loc[:, "annotated"] = h5file["annotated"][self.image_idx]
def _last_image(self):
""" Checks if image index is on the last index.
Helper method to check for the index of the last image in the h5 file.
Returns
-------
bool
Indicate if image_idx is the last index.
"""
return self.image_idx == self.n_images - 1
def _move_image_idx(self):
""" Move to different image.
Based on the key pressed, updates the image on the GUI.
The scheme is as follows:
------------------------------------------------------------
Keys | Action
------------------------------------------------------------
> <- , -> | Load previous or next image
> , , . | Jump 10 images backward or forward
------------------------------------------------------------
Every time the user moves from the image, the annotations
on the image is saved before loading the next image.
"""
# <- (left arrow) key
if self.key is keys.LEFTARROW:
self._save()
if self.image_idx == 0:
self.image_idx = self.n_images - 1
else:
self.image_idx -= 1
self._load()
# -> (right arrow) key
elif self.key is keys.RIGHTARROW:
self._save()
if self._last_image():
self.image_idx = 0
else:
self.image_idx += 1
self._load()
# . (period) key
elif self.key is keys.LESSTHAN:
self._save()
if self.image_idx - 10 < 0:
self.image_idx = self.n_images + self.image_idx - 10
else:
self.image_idx -= 10
self._load()
# , (comma) key
elif self.key is keys.GREATERTHAN:
self._save()
if self.image_idx + 10 > self.n_images - 1:
self.image_idx = self.image_idx + 10 - self.n_images
else:
self.image_idx += 10
self._load()
def _data(self):
""" Activates key bindings for annotated and save.
Creates additional key bindings for the program.
The bindings are as follows:
------------------------------------------------------------
Keys | Action
------------------------------------------------------------
> Ctrl-R | Mark frame as unannotated
> Ctrl-F | Mark frame as annotated
> Ctrl-S | Save
------------------------------------------------------------
"""
if self.key is keys.R:
self.skeleton["annotated"] = False
elif self.key is keys.F:
self.skeleton["annotated"] = True
elif self.key is keys.V:
if self.skeleton.loc[self.idx, ["x", "y"]].isnull()[0]:
self.skeleton.loc[self.idx, ["x", "y"]] = -1
else:
self.skeleton.loc[self.idx, ["x", "y"]] = np.nan
elif self.key in [keys.Q, keys.ESC]:
self._save()
print("Saved")
def _hotkeys(self):
""" Activates all key bindings.
Enables all the key functionalities described at the
start of the file.
"""
if self.key != keys.NONE:
self._wasd()
self._move_idx()
self._move_image_idx()
self._zoom()
self._data()
self._update_canvas()