-
Notifications
You must be signed in to change notification settings - Fork 0
/
flow_detect.py
71 lines (55 loc) · 2 KB
/
flow_detect.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
#!/usr/bin/env python3
"""DEPRECATED
"""
import argparse
import cv2
from flow import dump_flow, pick_flow, vis_flow
from detect import Detector
def draw_box_flow(frame, flow, detector):
bbox, label, score = detector(frame)
frame_rows = frame.shape[0]
frame_cols = frame.shape[1]
assert(frame.shape[2] == 3)
rows = flow.shape[0]
cols = flow.shape[1]
assert(flow.shape[2] == 2)
for i in range(rows):
for j in range(cols):
dx = flow[i, j, 0]
dy = flow[i, j, 1]
start = (j * frame_cols / cols + frame_cols / cols / 2,
i * frame_rows / rows + frame_rows / rows / 2)
start = tuple(map(int, start))
end = (start[0] + dx, start[1] + dy)
end = tuple(map(int, end))
cv2.line(frame, start, end, (0, 0, 255), 2)
for box in bbox:
cv2.rectangle(frame, tuple(box[0:2]), tuple(box[2:4]),
(0, 255, 0), 2)
return frame
def parseopt():
parser = argparse.ArgumentParser()
parser.add_argument("movie",
default="mpi_sinel_final_alley_1",
help="source movie dir to extract motion vectors")
parser.add_argument("--nobbox", "-n",
action="store_true", default=False,
help="dump occupancy enabled version")
parser.add_argument("--model",
choices=("ssd300", "ssd512"), default="ssd300")
parser.add_argument("--gpu", "-g", type=int, default=-1)
parser.add_argument("--pretrained_model", default="voc0712")
return parser.parse_args()
def main():
args = parseopt()
dump_flow(args.movie)
flow, header = pick_flow(args.movie)
detector = Detector(gpu=args.gpu)
draw_box_flow_func = \
lambda movie, flow: draw_box_flow(movie, flow, detector)
if args.nobbox:
vis_flow(args.movie, flow)
else:
vis_flow(args.movie, flow, draw=draw_box_flow_func)
if __name__ == "__main__":
main()