73 lines
2.7 KiB
Python
Executable File
73 lines
2.7 KiB
Python
Executable File
from PIL import Image
|
|
from util import *
|
|
import json
|
|
|
|
|
|
detector = FaceDetector('./weights/detection.onnx')
|
|
fer = HSEmotionRecognizer('./weights/emotion.onnx')
|
|
|
|
def detect_face(frame):
|
|
boxes = detector.detect(frame, (640, 640))
|
|
return boxes if boxes is not None and len(boxes) else None
|
|
|
|
|
|
def porcess_video(input_video_path):
|
|
stream = cv2.VideoCapture(input_video_path)
|
|
frame_width = int(stream.get(3))
|
|
frame_height = int(stream.get(4))
|
|
fps = stream.get(cv2.CAP_PROP_FPS)
|
|
|
|
idx_class = fer.idx_to_class
|
|
frame_counter = 0
|
|
emotion_reg_list = []
|
|
emotion_score_list = []
|
|
|
|
while True:
|
|
ret, frame = stream.read()
|
|
if ret:
|
|
frame_counter += 1
|
|
## 每30帧进行处理一次
|
|
if frame_counter % 60 == 0:
|
|
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
boxes = detect_face(image)
|
|
if boxes is not None:
|
|
for box in boxes.astype('int32'):
|
|
x1, y1, x2, y2 = box[:4]
|
|
face_image = image[y1:y2, x1:x2]
|
|
pil_image = Image.fromarray(face_image).convert('RGB')
|
|
pil_image = pil_image.resize((224, 224)) # Adjust size as per model requirements
|
|
emotion, scores,pred = fer.predict_emotions(np.array(pil_image), logits=False)
|
|
###### 做应用逻辑
|
|
print("frame_id:",frame_counter)
|
|
print("scores",scores[pred])
|
|
emotion_utf_8 = emotion.encode('utf-8').decode('utf-8')
|
|
print("emotion_utf_8",emotion_utf_8)
|
|
emotion_reg_list.append(emotion_utf_8)
|
|
emotion_score_list.append(float(scores[pred]))
|
|
print("idx_class",idx_class)
|
|
######
|
|
draw_emotion_bars(fer, frame, scores, (x2 + 10, y1), bar_height=15, width=100)
|
|
else:
|
|
print("退出")
|
|
break
|
|
|
|
stream.release()
|
|
|
|
|
|
data = {
|
|
"video_url":input_video_path,
|
|
"emotion_reg_list":emotion_reg_list,
|
|
"emotion_score_list":emotion_score_list
|
|
}
|
|
print("data:",data)
|
|
json_str = json.dumps(data)
|
|
print("json_str:",json_str)
|
|
return json_str,emotion_reg_list,emotion_score_list
|
|
|
|
if __name__ == '__main__':
|
|
#自己录制一段视频
|
|
video_path = "/Users/apple/Downloads/教学代码/111111.mp4"
|
|
porcess_video(video_path)
|
|
|
|
#结果如下
|
|
# json_str: {"video_url": "/Users/apple/Downloads/\u6559\u5b66\u4ee3\u7801/111111.mp4", "emotion_reg_list": ["NEUTRAL", "HAPPINESS"], "emotion_score_list": [0.5329781174659729, 0.43945103883743286]} |