Add tiny yolo

cvlib/object_detection.py

@@ -55,29 +55,37 @@ def draw_bbox(img, bbox, labels, confidence, colors=None, write_conf=False):
 
     return img
 
-def detect_common_objects(image, confidence=0.5, nms_thresh=0.3):
+def detect_common_objects(image, confidence=0.5, nms_thresh=0.3, model='yolov3'):
 
     Height, Width = image.shape[:2]
     scale = 0.00392
 
     global classes
     global dest_dir
 
-    config_file_name = 'yolov3.cfg'
-    config_file_abs_path = dest_dir + os.path.sep + config_file_name
+    if model == 'yolov3-tiny':
+        config_file_name = 'yolov3-tiny.cfg'
+        cfg_url = "https://github.com/pjreddie/darknet/raw/master/cfg/yolov3-tiny.cfg"
+        weights_file_name = 'yolov3-tiny.weights'
+        weights_url = 'https://pjreddie.com/media/files/yolov3-tiny.weights'
+        blob = cv2.dnn.blobFromImage(image, scale, (416,416), (0,0,0), True, crop=False)
+
 
-    weights_file_name = 'yolov3.weights'
+    else:
+        config_file_name = 'yolov3.cfg'
+        cfg_url = 'https://github.com/arunponnusamy/object-detection-opencv/raw/master/yolov3.cfg'
+        weights_file_name = 'yolov3.weights'
+        weights_url = 'https://pjreddie.com/media/files/yolov3.weights'
+        blob = cv2.dnn.blobFromImage(image, scale, (416,416), (0,0,0), True, crop=False)    
+
+    config_file_abs_path = dest_dir + os.path.sep + config_file_name
     weights_file_abs_path = dest_dir + os.path.sep + weights_file_name    
 
-    url = 'https://github.com/arunponnusamy/object-detection-opencv/raw/master/yolov3.cfg'
-
     if not os.path.exists(config_file_abs_path):
-        download_file(url=url, file_name=config_file_name, dest_dir=dest_dir)
-
-    url = 'https://pjreddie.com/media/files/yolov3.weights'
+        download_file(url=cfg_url, file_name=config_file_name, dest_dir=dest_dir)
 
     if not os.path.exists(weights_file_abs_path):
-        download_file(url=url, file_name=weights_file_name, dest_dir=dest_dir)    
+        download_file(url=weights_url, file_name=weights_file_name, dest_dir=dest_dir)    
 
     global initialize
     global net
@@ -87,8 +95,6 @@ def detect_common_objects(image, confidence=0.5, nms_thresh=0.3):
         net = cv2.dnn.readNet(weights_file_abs_path, config_file_abs_path)
         initialize = False
 
-    blob = cv2.dnn.blobFromImage(image, scale, (416,416), (0,0,0), True, crop=False)
-
     net.setInput(blob)
 
     outs = net.forward(get_output_layers(net))

examples/object_detection_video_yolov3_tiny.py

@@ -0,0 +1,47 @@
+# author: Arun Ponnusamy
+# website: https://www.arunponnusamy.com
+
+# object detection video example using tiny yolo model.
+# usage: python object_detection_video_yolov3_tiny.py /path/to/video
+
+# import necessary packages
+import cvlib as cv
+from cvlib.object_detection import draw_bbox
+import cv2
+import sys
+
+# open webcam
+video = cv2.VideoCapture(sys.argv[1])
+
+if not video.isOpened():
+    print("Could not open video")
+    exit()
+
+
+# loop through frames
+while video.isOpened():
+
+    # read frame from webcam 
+    status, frame = video.read()
+
+    if not status:
+        break
+
+    # apply object detection
+    bbox, label, conf = cv.detect_common_objects(frame, confidence=0.25, model='yolov3-tiny')
+
+    print(bbox, label, conf)
+
+    # draw bounding box over detected objects
+    out = draw_bbox(frame, bbox, label, conf, write_conf=True)
+
+    # display output
+    cv2.imshow("Real-time object detection", out)
+
+    # press "Q" to stop
+    if cv2.waitKey(1) & 0xFF == ord('q'):
+        break
+
+# release resources
+video.release()
+cv2.destroyAllWindows()        

examples/object_detection_webcam_yolov3_tiny.py

@@ -0,0 +1,46 @@
+# author: Arun Ponnusamy
+# website: https://www.arunponnusamy.com
+
+# object detection webcam example using tiny yolo
+# usage: python object_detection_webcam_yolov3_tiny.py
+
+# import necessary packages
+import cvlib as cv
+from cvlib.object_detection import draw_bbox
+import cv2
+
+# open webcam
+webcam = cv2.VideoCapture(0)
+
+if not webcam.isOpened():
+    print("Could not open webcam")
+    exit()
+
+
+# loop through frames
+while webcam.isOpened():
+
+    # read frame from webcam 
+    status, frame = webcam.read()
+
+    if not status:
+        break
+
+    # apply object detection
+    bbox, label, conf = cv.detect_common_objects(frame, confidence=0.25, model='yolov3-tiny')
+
+    print(bbox, label, conf)
+
+    # draw bounding box over detected objects
+    out = draw_bbox(frame, bbox, label, conf, write_conf=True)
+
+    # display output
+    cv2.imshow("Real-time object detection", out)
+
+    # press "Q" to stop
+    if cv2.waitKey(1) & 0xFF == ord('q'):
+        break
+
+# release resources
+webcam.release()
+cv2.destroyAllWindows()