-
Notifications
You must be signed in to change notification settings - Fork 2
/
model_test.py
57 lines (43 loc) · 2.03 KB
/
model_test.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
import pickle
import pandas as pd
import numpy as np
import cv2
import mediapipe as mp
def img_model_test(image, model):
mp_drawing = mp.solutions.drawing_utils # Drawing helpers.
mp_holistic = mp.solutions.holistic # Mediapipe Solutions.
with mp_holistic.Holistic(min_detection_confidence=0.5, min_tracking_confidence=0.5) as holistic:
# Recolor Feed
image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
image.flags.writeable = False
# Make Detections
results = holistic.process(image)
# Recolor image back to BGR for rendering
image.flags.writeable = True
image = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
# Pose Detections
mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_holistic.POSE_CONNECTIONS,
mp_drawing.DrawingSpec(color=(245,117,66), thickness=2, circle_radius=4),
mp_drawing.DrawingSpec(color=(245,66,230), thickness=2, circle_radius=2)
)
# Extract Pose landmarks
pose = results.pose_landmarks.landmark
pose_row = list(np.array([[landmark.x, landmark.y, landmark.z, landmark.visibility] for landmark in pose]).flatten())
# Concate rows
row = pose_row
# Make Detections
X = pd.DataFrame([row])
body_language_class = model.predict(X)[0]
body_language_prob = model.predict_proba(X)[0]
# print(f'class: {body_language_class}, prob: {body_language_prob}')
return body_language_class, body_language_prob
if __name__ == '__main__':
img = cv2.imread('./resource/imgs/test1.png')
model_weights = './model_weights/weights_body_language.pkl'
# Load Model.
with open(model_weights, 'rb') as f:
model = pickle.load(f)
# Input image to test model predict.
predict_class = img_model_test(image=img, model=model)[0]
class_probbile = img_model_test(image=img, model=model)[1]
print(f'\nclass: {predict_class}, acc: {class_probbile}\n')