Update utils/emotion_detection_brainai.py

utils/emotion_detection_brainai.py

@@ -12,8 +12,15 @@ class EmotionModel:
     def __init__(self):
         self.face_compiled_model, self.face_input_layer, self.face_output_layer = self.load_model('face-detection-adas-0001')
         self.emotion_compiled_model, self.emotion_input_layer, self.emotion_output_layer = self.load_model('emotions-recognition-retail-0003')
-        
 
+        self.emotions = {
+            0:"neutral",
+            1:"happy",
+            2:"sad",
+            3:"surprise",
+            4:"anger"
+        }
+        
     def load_model(self, model_name):
         model_path = "models/" + model_name + ".xml"
         core = ov.Core()
@@ -32,6 +39,11 @@ class EmotionModel:
         
         return input_img
 
+    def detect_faces(self, img, conf = 0.5)
+        result_face = self.face_compiled_model([input_img])[self.face_output_layer]
+        boxes = self.post_process_face(result_face, uploaded_img_cv, conf=0.5)
+
+        return boxes
 
     def post_process_face(self, result_face, img, conf=0.5):
         boxes = []
@@ -49,19 +61,19 @@ class EmotionModel:
         
         return boxes
 
+    def detect_emotions(self, img, boxes)
+        for box in boxes:
+            xmin, ymin, xmax, ymax = box
+            emotion_input = uploaded_img_cv[ymin:ymax,xmin:xmax]
+            input_img = self.preprocess(emotion_input, self.emotion_input_layer)       
+            result_emotion = self.emotion_compiled_model([input_img])[self.emotion_output_layer]
+            self.post_process_emotion(result_emotion, uploaded_img_cv, box)
+
     def post_process_emotion(self, result_emotion, img, face_position):
         
-        emotions = {
-            0:"neutral",
-            1:"happy",
-            2:"sad",
-            3:"surprise",
-            4:"anger"
-        }
-                    
         predictions = result_emotion[0,:,0,0]   
         topresult_index = np.argmax(predictions)
-        emotion = emotions[topresult_index]
+        emotion = self.emotions[topresult_index]
 
         font_size = img.shape[0]/1000
         font_thickness = int(img.shape[0]/500)
@@ -82,19 +94,9 @@ class EmotionModel:
             uploaded_img = PIL.Image.open(img)
             uploaded_img_cv =  np.array(uploaded_img) 
 
-        input_img = self.preprocess(uploaded_img_cv, self.face_input_layer)
-        result_face = self.face_compiled_model([input_img])[self.face_output_layer]
-        boxes = self.post_process_face(result_face, uploaded_img_cv, conf=0.5)
-
-        if boxes is not None:
-
-            for box in boxes:
-                xmin, ymin, xmax, ymax = box
-                emotion_input = uploaded_img_cv[ymin:ymax,xmin:xmax]
-                input_img = self.preprocess(emotion_input, self.emotion_input_layer)       
-                result_emotion = self.emotion_compiled_model([input_img])[self.emotion_output_layer]
-                self.post_process_emotion(result_emotion, uploaded_img_cv, box)
-
+        boxes = detect_faces(uploaded_img_cv, conf = 0.5)
+        detect_emotions(uploaded_img_cv)
+        
         return uploaded_img_cv
 
 emotion_model = EmotionModel()