This repository has been archived by the owner on Apr 11, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 0
/
find_bb.py
147 lines (114 loc) · 4.81 KB
/
find_bb.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
#!/usr/bin/env python3
# Copyright 2018 Google LLC
#
# 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 cv2
import sys
import os
import numpy as np
import bbox_writer
import argparse
import drawing_utils
description_text="""\
Use this script to draw initial bounding boxes on videos.
This script is meant to work hand in hand with tracking.py. To label a new
video, you should first annotate bounding boxes on the video's first frame with
this script, and then load the video into tracking.py to track the bounding
boxes.
To label an object, you should first pick the class that the object belongs to,
by pressing any (lowercase) alphanumeric character on your keyboard. This will
set the class to be that single character, e.g. 'v', 'a', 's', 'u', etc. After
picking a class, you can use your mouse to click on two opposing corners of a
box around an object. After you click the second point, the object will be
labeled with the class you specified earlier. At this point, you can change the
class, or draw another box with the same class.
Once you're satisfied with the boxes you've drawn, you can exit the program by
pressing 'q'. This also means that 'q' is not a valid object class.
If you would like to clear the first point you selected, you can press Escape.
If you would like to remove a bounding box you've drawn already, you should
simply restart the program, which will clear all bounding boxes.
Finally, you should try to make bounding boxes drawn here as tight as possible.
Any scaling of bounding boxes that needs to be performed can be done in
tracking.py.
"""
parser = argparse.ArgumentParser(
description=description_text,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("filename", type=argparse.FileType("r"),
help="path to video file")
args = parser.parse_args()
window = 'draw on me!'
points = []
bboxes = []
classes = []
current_class = None
WINDOW_SCALE = 0.5
def show_scaled(window, frame, sf=WINDOW_SCALE):
cv2.imshow(window, cv2.resize(frame, (0, 0), fx=sf, fy=sf))
# Make sure all bboxes are given as top left (x, y), and (dx, dy). Sometimes
# they may be specified by a different corner, so we need to reverse that.
def standardize_bbox(bbox):
p0 = bbox[:2]
p1 = p0 + bbox[2:]
min_x = min(p0[0], p1[0])
max_x = max(p0[0], p1[0])
min_y = min(p0[1], p1[1])
max_y = max(p0[1], p1[1])
ret = np.array([min_x, min_y, max_x - min_x, max_y - min_y])
print("Standardized %s to %s" % (bbox, ret))
return ret
def mouse_callback(event, x, y, flags, params):
im = params.copy()
h, w, c = im.shape
x = int(x / WINDOW_SCALE)
y = int(y / WINDOW_SCALE)
if event == cv2.EVENT_LBUTTONDOWN:
points.append(np.array([x, y]))
if len(points) == 1: # Still drawing a rectangle
cv2.rectangle(im, tuple(points[0]), (x, y), (255, 255, 0), 1, 1)
# If the mouse is moved, draw crosshairs
cv2.line(im, (x, 0), (x, h), (255, 0, 0))
cv2.line(im, (0, y), (w, y), (255, 0, 0))
if len(points) == 2: # We've got a rectangle
bbox = np.array([points[0], points[1] - points[0]]).reshape(-1)
bbox = standardize_bbox(bbox)
cls = str(current_class)
bboxes.append(bbox)
classes.append(cls)
points.clear()
# Write the bboxes out to file.
# We get the filename itself, not the full path.
filename = args.filename.name
bbox_filename = bbox_writer.get_bbox_filename(filename)
bbox_path = os.path.join(os.path.dirname(filename), bbox_filename)
bbox_writer.write_bboxes(bboxes, classes, bbox_path)
print("Wrote bboxes to file:", bbox_path)
drawing_utils.draw_bboxes(im, bboxes, classes)
cv2.putText(im, "Current class: %s" % current_class, (100,20),
cv2.FONT_HERSHEY_SIMPLEX, 0.75, (50, 170, 50), 2)
show_scaled(window, im)
if __name__ == "__main__":
vid = cv2.VideoCapture(args.filename.name)
cv2.namedWindow(window)
if vid.isOpened():
ret, frame = vid.read()
show_scaled(window, frame)
cv2.setMouseCallback(window, mouse_callback, param=frame)
while True:
key = cv2.waitKey(1) & 0xFF
if key == ord('q'):
break
elif key == 27: # esc
points.clear()
elif key > -1 and key < 128 and chr(key).isalnum():
current_class = chr(key)