-
Notifications
You must be signed in to change notification settings - Fork 5
/
predict.py
149 lines (102 loc) · 4.31 KB
/
predict.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
from __future__ import generators
import sys
import caffe
import matplotlib
import numpy as np
import h5py
import glob
import re
import argparse
from collections import defaultdict
from frameIO import *
from frameset import *
from extract_frames import *
NUMBER_OF_LABELS = 37
def read_means(means_file_name):
with io.open(means_file_name, "r") as f:
s = f.read()
return json.loads(s)
def avg_net_results(probs):
return reduce(lambda x, y: x + y, probs) / 2
def get_predictions(video_file_name, means_file_name="videos/means.json"):
frames = get_frames(video_file_name)
means = read_means(means_file_name)
frameSets = preprocess_frames(frames, means)
prob_result = forward_net_multi(frameSets[0].frames, "../face-vid-nets/one-vs-all/deploy.prototxt", "../face-vid-nets/snapshots/ONE-*_iter_2500.caffemodel")
# prob_result = avg_net_results([
# forward_net_single(frameSets[1], "nets/an-finetune/deploy.prototxt", "nets/ANF_iter_50000.caffemodel"),
# forward_net_multi(frameSets[1], "face-vid-nets/one-vs-all/deploy.prototxt", "face-vid-nets/snapshots/ONE-*_iter_10000.caffemodel")
# ])
return prob_result
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--video', type=str, required=True)
parser.add_argument('--means', type=str, required=True)
args = parser.parse_args()
video_file_name = args.video
means_file_name = args.means
frames = get_frames(video_file_name)
means = read_means(means_file_name)
frameSets = preprocess_frames(frames, means)
# prob_result = avg_net_results([
# forward_net_single(frameSets[0].frames, "nets/an-finetune/deploy.prototxt", "nets/ANF_iter_50000.caffemodel"),
# forward_net_multi(frameSets[0].frames, "face-vid-nets/one-vs-all/deploy.prototxt", "face-vid-nets/snapshots/ONE-*_iter_10000.caffemodel")
# ])
prob_result = forward_net_multi(frameSets[0].frames, "face-vid-nets/one-vs-all/deploy.prototxt", "face-vid-nets/snapshots/ONE-*_iter_10000.caffemodel")
print prob_result
label_result = prob_result > THRESHOLD
print label_result
def preprocess_frames(frames, means):
processId = id_generator()
face_cache = {}
labels = np.zeros((1,1,1,1))
# Extract frames
frameSet = FrameSet(frames, "framesOriginal", processId, labels)
frameSets = split_grayscale_BGR(frameSet)
frameSets = detect_faces_and_mask_surroundings(frameSets, face_cache)
frameSets = induce_flows(frameSets)
frameSets = filter_framesets_out_by_stream_name(frameSets, "grayscale")
# frameSets = filter_frames_with_labels(frameSets)
frameSets = resize_frames(frameSets, 227, 227)
# frameSets = accumulate_means(frameSets, means, layer_counts)
frameSets = transform_to_caffe_format(frameSets)
# Finalize (e.g. mean substraction)
frameSets = substract_means(frameSets, means)
frameSets = set_mask_to_zero(frameSets)
frameSets = cross_flows(frameSets)
frameSets = list(frameSets)
frameSetsBGR = filter(lambda f: f.streamName == "BGR", frameSets)
frameSetsFlow = filter(lambda f: f.isFlow(), frameSets)
frameSetsFlow = list(normalize_frames(frameSetsFlow))
return frameSetsBGR + frameSetsFlow
def forward_net_single(data, proto, model):
net = caffe.Net(proto, model, caffe.TEST)
caffe.set_mode_gpu()
print "Starting to forward data..."
out = net.forward_all(data=data)
print "Prediction finished"
return out['prob']
def forward_net_multi(data, proto, model_pattern):
models = glob.glob(model_pattern)
# output = np.zeros((data.shape[0], len(models)))
output = np.zeros((data.shape[0], NUMBER_OF_LABELS))
for model in models:
i = int(re.match(model_pattern.replace("*", "(\d+)"), model).group(1))
net = caffe.Net(proto, model, caffe.TEST)
caffe.set_mode_gpu()
print "(%i) Starting to forward data..." % i
out = net.forward_all(data=data)
print "(%i) Prediction finished" % i
prob_labels = out["prob"]
output[:, i] = prob_labels[:, 1]
return output
def get_labels():
labels = {}
with open("videos/labelmapping.txt","r") as f:
next(f) # skip header
for line in f:
value, key = line.split()
labels[int(key)] = value
return labels
if __name__ == "__main__":
main()