60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
import sys
|
|
|
|
from algorithm.detect_emotion.rmn import RMN
|
|
from PIL import Image
|
|
import cv2
|
|
import matplotlib.pyplot as plt # plt 用于显示图片
|
|
from read_data import LoadImages, LoadStreams
|
|
import torch
|
|
import time
|
|
import torch.backends.cudnn as cudnn
|
|
|
|
class Emotion_Detection():
|
|
def __init__(self,video_path=None, model = None):
|
|
|
|
self.frame = [None]
|
|
|
|
if video_path is not None:
|
|
self.video_name = video_path
|
|
else:
|
|
self.video_name = 'vid2.mp4' # A default video file
|
|
|
|
self.dataset = LoadImages(self.video_name)
|
|
self.face_detector = model
|
|
self.emotion_model = RMN(face_detector = self.face_detector)
|
|
|
|
def use_webcam(self, source):
|
|
# self.dataset.release() # Release any existing video capture
|
|
# self.cap = cv2.VideoCapture(0) # Open default webcam
|
|
# print('use_webcam')
|
|
source = source
|
|
cudnn.benchmark = True
|
|
self.dataset = LoadStreams(source)
|
|
|
|
def get_frame(self):
|
|
|
|
for im0s in self.dataset:
|
|
|
|
if self.dataset.mode == 'stream':
|
|
img = im0s[0].copy()
|
|
else:
|
|
img = im0s.copy()
|
|
|
|
|
|
results = self.emotion_model.detect_emotion_for_single_frame(img)
|
|
|
|
keyword_to_remove = 'proba_list'
|
|
image = self.emotion_model.draw(img, results)
|
|
|
|
for dictionary in results:
|
|
if keyword_to_remove in dictionary:
|
|
del dictionary[keyword_to_remove]
|
|
# print(results)
|
|
|
|
|
|
ret, jpeg = cv2.imencode(".jpg", image)
|
|
return jpeg.tobytes(), ''
|
|
|
|
|
|
# x.Emotion_result(picpath="666666.png")
|