Updated main.py

 src/model.py

Signed-off-by: Mpho Mphego <mpho112@gmail.com>

Author: mmphego

main.py

@@ -141,7 +141,10 @@ def main(args):
     logger.info(f"Total time taken to load all the models: {model_load_time:.2f} secs.")
 
     for frame in video_feed.next_frame():
-        predict_end_time, _, face_bboxes = face_detection.predict(frame, draw=True)
+        if args.debug:
+            video_feed.show(video_feed.resize(frame))
+
+        predict_end_time, face_bboxes = face_detection.predict(frame, draw=True)
         text = f"Face Detection Inference time: {predict_end_time:.3f} s"
         face_detection.add_text(text, frame, (15, video_feed.source_height - 80))
 
@@ -164,15 +167,12 @@ def main(args):
                 if face_height < 20 or face_width < 20:
                     continue
 
-                predict_end_time, _, landmarks_bboxes = facial_landmarks.predict(face, draw=True)
+                predict_end_time, eyes_coords = facial_landmarks.predict(face, draw=True)
                 text = f"Facial Landmarks Est. Inference time: {predict_end_time:.3f} s"
                 facial_landmarks.add_text(
                     text, frame, (15, video_feed.source_height - 60)
                 )
 
-        if args.debug:
-            video_feed.show(video_feed.resize(frame))
-
     video_feed.close()
 
 

src/model.py

@@ -109,7 +109,7 @@ def predict(self, image, request_id=0, draw=False):
             predict_end_time = float(time.time() - predict_start_time) * 1000
             if draw:
                 bbox, _ = self.preprocess_output(pred_result, image, show_bbox=draw)
-            return (predict_end_time, pred_result, bbox)
+            return (predict_end_time, bbox)
 
     @abstractmethod
     def preprocess_output(self, inference_results, image, show_bbox=False):
@@ -257,14 +257,15 @@ def preprocess_output(self, inference_results, image, show_bbox=False):
         }
         if show_bbox:
             self.draw_output(image, eyes_coords)
-        return eyes_coords
+        return eyes_coords, image
 
     @staticmethod
     def draw_output(image, eyes_coords, radius=10, color=(0, 0, 255), thickness=2):
         """Draw a circle around ROI"""
         for eye, coords in eyes_coords.items():
             cv2.circle(image, (coords[0], coords[1]), radius, color, thickness)
 
+
 class Head_Pose_Estimation(Base):
     """Class for the Head Pose Estimation Model."""