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")