Merge pull request sthanhng#2 from sthanhng/livevideo

YOLOFace - Livevideo

YOLO.py

@@ -147,33 +147,54 @@ def detect_image(self, image):
             for thk in range(thickness):
                 draw.rectangle(
                     [left + thk, top + thk, right - thk, bottom - thk],
-                    outline=self.colors[c])
+                    outline=(51, 178, 255))
             del draw
 
         end_time = timer()
         print('[i] ==> Processing time: {:.2f}ms'.format((end_time -
                                                           start_time) * 1000))
-        return image
+        return image, out_boxes
 
     def close_session(self):
         self.sess.close()
 
 
+def detect_img(yolo):
+    while True:
+        img = input('[i] ==> Input image filename: ')
+        try:
+            image = Image.open(img)
+        except:
+            print('[!] ==> Open Error! Try again!')
+            continue
+        else:
+            res_image, _ = yolo.detect_image(image)
+            res_image.show()
+
+    yolo.close_session()
+
+
 def detect_video(model, video_path=None, output=None):
-    vid = cv2.VideoCapture(video_path)
+    if video_path == 'stream':
+        vid = cv2.VideoCapture(0)
+    else:
+        vid = cv2.VideoCapture(video_path)
+
     if not vid.isOpened():
         raise IOError("Couldn't open webcam or video")
 
     # The video format and fps
-    video_fourcc = int(vid.get(cv2.CAP_PROP_FOURCC))
+    # video_fourcc = int(vid.get(cv2.CAP_PROP_FOURCC))
+    video_fourcc = cv2.VideoWriter_fourcc('M', 'G', 'P', 'G')
     video_fps = vid.get(cv2.CAP_PROP_FPS)
 
     # The size of the frames to write
     video_size = (int(vid.get(cv2.CAP_PROP_FRAME_WIDTH)),
                   int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT)))
     isOutput = True if output != "" else False
     if isOutput:
-        out = cv2.VideoWriter(output, video_fourcc, video_fps, video_size)
+        output_fn = 'output_video.avi'
+        out = cv2.VideoWriter(os.path.join(output, output_fn), video_fourcc, video_fps, video_size)
 
     accum_time = 0
     curr_fps = 0
@@ -184,7 +205,7 @@ def detect_video(model, video_path=None, output=None):
         ret, frame = vid.read()
         if ret:
             image = Image.fromarray(frame)
-            image = model.detect_image(image)
+            image, faces = model.detect_image(image)
             result = np.asarray(image)
 
             curr_time = timer()
@@ -194,10 +215,22 @@ def detect_video(model, video_path=None, output=None):
             curr_fps = curr_fps + 1
             if accum_time > 1:
                 accum_time = accum_time - 1
-                fps = 'FPS: {}'.format(curr_fps)
+                fps = curr_fps
                 curr_fps = 0
-            cv2.putText(result, fps, (20, 30), cv2.FONT_HERSHEY_SIMPLEX,
-                        0.6, (0, 255, 0), 2)
+
+            # Initialize the set of information we'll displaying on the frame
+            info = [
+                ('FPS', '{}'.format(fps)),
+                ('Faces detected', '{}'.format(len(faces)))
+            ]
+
+            cv2.rectangle(result, (5, 5), (120, 50), (0, 0, 0), cv2.FILLED)
+
+            for (i, (txt, val)) in enumerate(info):
+                text = '{}: {}'.format(txt, val)
+                cv2.putText(result, text, (10, (i * 20) + 20),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.3, (10, 175, 0), 1)
+
             cv2.namedWindow("face", cv2.WINDOW_NORMAL)
             cv2.imshow("face", result)
             if isOutput:

yoloface_gpu.py

@@ -1,7 +1,7 @@
 import argparse
 
 from PIL import Image
-from YOLO import YOLO, detect_video
+from YOLO import YOLO, detect_video, detect_img
 
 
 #####################################################################
@@ -23,27 +23,12 @@ def get_args():
                         help='image detection mode')
     parser.add_argument('--video', type=str, default='samples/subway.mp4',
                         help='path to the video')
-    parser.add_argument('--output', type=str,
-                        default='', help='image/video output path')
+    parser.add_argument('--output', type=str, default='outputs/',
+                        help='image/video output path')
     args = parser.parse_args()
     return args
 
 
-def detect_img(yolo):
-    while True:
-        img = input('[i] ==> Input image filename: ')
-        try:
-            image = Image.open(img)
-        except:
-            print('[!] ==> Open Error! Try again!')
-            continue
-        else:
-            res_image = yolo.detect_image(image)
-            res_image.show()
-
-    yolo.close_session()
-
-
 def _main():
     # Get the arguments
     args = get_args()
@@ -55,7 +40,7 @@ def _main():
     else:
         print('[i] ==> Video detection mode\n')
         # Call the detect_video method here
-        detect_video(YOLO(args), args.video)
+        detect_video(YOLO(args), args.video, args.output)
 
     print('Well done!!!')