forked from smellslikeml/ActionAI
-
Notifications
You must be signed in to change notification settings - Fork 0
/
inference.py
32 lines (23 loc) · 913 Bytes
/
inference.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
if __name__ == '__main__':
import cv2
import pickle
import argparse
import importlib
from transformer import PoseExtractor
parser = argparse.ArgumentParser(description='Run inference on webcam video')
parser.add_argument('--config', type=str, default='conf',
help="name of config .py file inside config/ directory, default: 'conf'")
args = parser.parse_args()
config = importlib.import_module('config.' + args.config)
model = pickle.load(open(config.classifier_model, 'rb'))
extractor = PoseExtractor()
cap = cv2.VideoCapture(config.stream)
while True:
ret, image = cap.read()
sample = extractor.transform([image])
prediction = model.predict(sample.reshape(1, -1))
print(prediction[0])
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()