added face detection

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

Author: mmphego

main.py

@@ -136,7 +136,15 @@ def main(args):
 
     video_feed = InputFeeder(input_file=args.input)
 
+    # Add source width and height for face detection.
+    face_detection._init_image_w = video_feed.source_width
+    face_detection._init_image_h = video_feed.source_height
+
     for frame in video_feed.next_frame():
+        predict_end_time, pred_result = face_detection.predict(frame,draw=True)
+        text = f"Inference time: {predict_end_time:.2f}ms"
+        face_detection.add_text(text, frame, (15, face_detection._init_image_h - 50))
+
         if args.debug:
             video_feed.show(frame)
 

src/input_feeder.py

@@ -78,8 +78,11 @@ def progress_bar(self):
             self._progress_bar = tqdm(total=int(self.video_len - self.fps + 1))
         return self._progress_bar
 
+    def resize(self,frame):
+        return cv2.resize(frame, (self.source_width - 200, self.source_height - 200))
+
     def show(self, frame, frame_name="video"):
-        cv2.imshow(frame_name, frame)
+        cv2.imshow(frame_name, self.resize(frame))
 
     def write_video(self, output_path=".", filename="output_video.mp4"):
         out_video = cv2.VideoWriter(
@@ -100,15 +103,14 @@ def next_frame(self, quit_key="q"):
                 flag, frame = self.cap.read()
 
             if not flag:
-                    break
+                break
             yield frame
 
             key = cv2.waitKey(1) & 0xFF
             # if `quit_key` was pressed, break from the loop
             if key == ord(quit_key):
                 break
 
-
     def close(self):
         """Closes the VideoCapture."""
         if "image" in self._input_type:

src/model.py

@@ -92,21 +92,30 @@ def predict(self, image, request_id=0, draw=False):
         )
         status = self.exec_network.requests[request_id].wait(-1)
         if status == 0:
+            predict_start_time = time.time()
             pred_result = self.exec_network.requests[request_id].outputs[
                 self.output_name
             ]
-            return pred_result
+            predict_end_time = (time.time() - predict_start_time) * 1000
+            if draw:
+                self.preprocess_output(pred_result, image, show_bbox=draw)
+            return (predict_end_time, pred_result)
 
     @abc.abstractmethod
-    def preprocess_output(self, inference_results, image):
+    def preprocess_output(self, inference_results, image, show_bbox=False):
         """Draw bounding boxes onto the frame."""
         pass
 
     @staticmethod
     @abc.abstractmethod
-    def draw_output(coords, image):
+    def draw_output(image, ):
         pass
 
+    def add_text(self, text, frame, position, font_size=0.75, color=(255, 255, 255)):
+        cv2.putText(
+            frame, text, position, cv2.FONT_HERSHEY_COMPLEX, font_size, color, 1,
+        )
+
     def preprocess_input(self, image):
         """Helper function for processing frame"""
         p_frame = cv2.resize(image, (self.input_shape[3], self.input_shape[2]))
@@ -122,7 +131,7 @@ class Face_Detection(Base):
     def __init__(self, model_name, device="CPU", threshold=0.60, extensions=None):
         super().__init__(model_name, device="CPU", threshold=0.60, extensions=None)
 
-    def preprocess_output(self, inference_results, image):
+    def preprocess_output(self, inference_results, image, show_bbox=False):
         """Draw bounding boxes onto the frame."""
         if not (self._init_image_w and self._init_image_h):
             raise RuntimeError("Initial image width and height cannot be None.")
@@ -136,10 +145,13 @@ def preprocess_output(self, inference_results, image):
                 xmax = int(box[5] * self._init_image_w)
                 ymax = int(box[6] * self._init_image_h)
                 coords.append((xmin, ymin, xmax, ymax))
+                if show_bbox:
+                    self.draw_output(image, xmin, ymin, xmax, ymax)
         return coords, image
 
-    def draw_output(coords, image):
-        label = "Person"
+    @staticmethod
+    def draw_output(image, xmin, ymin, xmax, ymax):
+        label = "Person's Face"
         bbox_color = (0, 255, 0)
         padding_size = (0.05, 0.25)
         text_color = (255, 255, 255)