Add object detection code

cvlib/__init__.py

@@ -1,6 +1,7 @@
 # author: Arun Ponnusamy
 # website: https://www.arunponnusamy.com
 
-__version__ = "0.1.3"
+__version__ = "0.1.6"
 
-from .face_detection import detect_face 
+from .face_detection import detect_face
+from .object_detection import detect_common_objects

cvlib/object_detection.py

@@ -0,0 +1,136 @@
+import cv2
+import os
+import numpy as np
+from .utils import download_file
+
+initialize = True
+net = None
+dest_dir = os.path.expanduser('~') + os.path.sep + '.cvlib' + os.path.sep + 'object_detection' + os.path.sep + 'yolo' + os.path.sep + 'yolov3'
+classes = None
+COLORS = np.random.uniform(0, 255, size=(80, 3))
+
+def populate_class_labels():
+
+    class_file_name = 'yolov3_classes.txt'
+    class_file_abs_path = dest_dir + os.path.sep + class_file_name
+    url = 'https://github.com/arunponnusamy/object-detection-opencv/raw/master/yolov3.txt'
+    if not os.path.exists(class_file_abs_path):
+        download_file(url=url, file_name=class_file_name, dest_dir=dest_dir)
+    f = open(class_file_abs_path, 'r')
+    classes = [line.strip() for line in f.readlines()]
+
+    return classes
+
+
+def get_output_layers(net):
+
+    layer_names = net.getLayerNames()
+
+    output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
+
+    return output_layers
+
+
+def draw_bbox(img, bbox, labels, confidence, colors=None, write_conf=False):
+
+    global COLORS
+    global classes
+
+    if classes is None:
+        classes = populate_class_labels()
+
+    for i, label in enumerate(labels):
+
+        if colors is None:
+            color = COLORS[classes.index(label)]            
+        else:
+            color = colors[classes.index(label)]
+
+        if write_conf:
+            label += ' ' + str(format(confidence[i] * 100, '.2f')) + '%'
+
+        cv2.rectangle(img, (bbox[i][0],bbox[i][1]), (bbox[i][2],bbox[i][3]), color, 2)
+
+        cv2.putText(img, label, (bbox[i][0],bbox[i][1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+    return img
+
+def detect_common_objects(image):
+
+    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
+
+    weights_file_name = 'yolov3.weights'
+    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'
+
+    if not os.path.exists(weights_file_abs_path):
+        download_file(url=url, file_name=weights_file_name, dest_dir=dest_dir)    
+
+    global initialize
+    global net
+
+    if initialize:
+        classes = populate_class_labels()
+        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))
+
+    class_ids = []
+    confidences = []
+    boxes = []
+    conf_threshold = 0.5
+    nms_threshold = 0.4
+
+    for out in outs:
+        for detection in out:
+            scores = detection[5:]
+            class_id = np.argmax(scores)
+            confidence = scores[class_id]
+            if confidence > 0.5:
+                center_x = int(detection[0] * Width)
+                center_y = int(detection[1] * Height)
+                w = int(detection[2] * Width)
+                h = int(detection[3] * Height)
+                x = center_x - w / 2
+                y = center_y - h / 2
+                class_ids.append(class_id)
+                confidences.append(float(confidence))
+                boxes.append([x, y, w, h])
+
+
+    indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)
+
+    bbox = []
+    label = []
+    conf = []
+
+    for i in indices:
+        i = i[0]
+        box = boxes[i]
+        x = box[0]
+        y = box[1]
+        w = box[2]
+        h = box[3]
+        bbox.append([round(x), round(y), round(x+w), round(y+h)])
+        label.append(str(classes[class_ids[i]]))
+        conf.append(confidences[i])
+
+    return bbox, label, conf

cvlib/utils.py

@@ -0,0 +1,43 @@
+import requests
+import progressbar as pb
+import os
+
+def download_file(url, file_name, dest_dir):
+
+    if not os.path.exists(dest_dir):
+        os.makedirs(dest_dir)
+
+    full_path_to_file = dest_dir + os.path.sep + file_name
+
+    if os.path.exists(dest_dir + os.path.sep + file_name):
+        return full_path_to_file
+
+    print("Downloading " + file_name + " from " + url)
+
+    try:
+        r = requests.get(url, allow_redirects=True, stream=True)
+    except:
+        print("Could not establish connection. Download failed")
+        return None
+
+    file_size = int(r.headers['Content-Length'])
+    chunk_size = 1024
+    num_bars = round(file_size / chunk_size)
+
+    bar = pb.ProgressBar(maxval=num_bars).start()
+
+    if r.status_code != requests.codes.ok:
+        print("Error occurred while downloading file")
+        return None
+
+    count = 0
+
+    with open(full_path_to_file, 'wb') as file:
+        for chunk in  r.iter_content(chunk_size=chunk_size):
+            file.write(chunk)
+            bar.update(count)
+            count +=1
+
+    return full_path_to_file
+
+

setup.py

@@ -1,7 +1,7 @@
 from setuptools import setup
 
 setup(name='cvlib',
-      version='0.1.5',
+      version='0.1.6',
       description='A high level, easy to use, open source computer vision library for python',
       url='https://github.com/arunponnusamy/cvlib.git',
       author='Arun Ponnusamy',