pipeline impl

Server_Impl/pipeline.py

+import numpy as np
+import argparse
+import cv2
+import os
+import time
+from model.config import config
+
+
+def extract_boxes_confidences_classids(outputs, confidence, width, height):
+    boxes = []
+    confidences = []
+    classIDs = []
+
+    for output in outputs:
+        for detection in output:            
+            # Extract the scores, classid, and the confidence of the prediction
+            scores = detection[5:]
+            classID = np.argmax(scores)
+            conf = scores[classID]
+            
+            # Consider only the predictions that are above the confidence threshold
+            if conf > confidence:
+                # Scale the bounding box back to the size of the image
+                box = detection[0:4] * np.array([width, height, width, height])
+                centerX, centerY, w, h = box.astype('int')
+
+                # Use the center coordinates, width and height to get the coordinates of the top left corner
+                x = int(centerX - (w / 2))
+                y = int(centerY - (h / 2))
+
+                boxes.append([x, y, int(w), int(h)])
+                confidences.append(float(conf))
+                classIDs.append(classID)
+
+    return boxes, confidences, classIDs
+
+
+def draw_bounding_boxes(image,labels, boxes, confidences, classIDs, idxs, colors):
+    if len(idxs) > 0:
+        for i in idxs.flatten():
+            # extract bounding box coordinates
+            x, y = boxes[i][0], boxes[i][1]
+            w, h = boxes[i][2], boxes[i][3]
+
+            # draw the bounding box and label on the image
+            color = [int(c) for c in colors[classIDs[i]]]
+            cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
+            text = "{}: {:.4f}".format(labels[classIDs[i]], confidences[i])
+            cv2.putText(image, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+    return image
+
+
+def make_prediction(net, layer_names, labels, image, confidence, threshold):
+    height, width = image.shape[:2]
+    
+    # Create a blob and pass it through the model
+    blob = cv2.dnn.blobFromImage(image, 1 / 255.0, (416, 416), swapRB=True, crop=False)
+    net.setInput(blob)
+    outputs = net.forward(layer_names)
+
+    # Extract bounding boxes, confidences and classIDs
+    boxes, confidences, classIDs = extract_boxes_confidences_classids(outputs, confidence, width, height)
+
+    # Apply Non-Max Suppression
+    idxs = cv2.dnn.NMSBoxes(boxes, confidences, confidence, threshold)
+
+    return boxes, confidences, classIDs, idxs
+
+
+def pipeline(input_file,output_path):
+    # Get the labels
+    labels = open(config.labels).read().strip().split('\n')
+    # Create a list of colors for the labels
+    colors = np.random.randint(0, 255, size=(len(config.labels), 3), dtype='uint8')
+    # Load weights using OpenCV
+    net = cv2.dnn.readNetFromDarknet(config.config, config.weights)
+    # Get the ouput layer names
+    layer_names = net.getLayerNames()
+    layer_names = [layer_names[i- 1] for i in net.getUnconnectedOutLayers()]
+
+    cap = cv2.VideoCapture(input_file)
+
+    # We need to set resolutions.
+    # so, convert them from float to integer.
+    frame_width = int(cap.get(3))
+    frame_height = int(cap.get(4))
+    size = (frame_width, frame_height)
+    result = cv2.VideoWriter(output_path+'output.avi', 
+                         cv2.VideoWriter_fourcc(*'MJPG'),
+                         10, size)
+    if (cap.isOpened()== False):
+        # Give a error message
+        # Read the entire file until it is completed
+        print("error")
+    while(cap.isOpened()):
+    # Capture each frame
+        ret, image = cap.read()
+        if ret == True:
+            # Display the resulting frame
+            boxes, confidences, classIDs, idxs = make_prediction(net, layer_names, labels, image, config.confidence, config.threshold)
+            image = draw_bounding_boxes(image,labels, boxes, confidences, classIDs, idxs, colors)
+            result.write(image)
+            # cv2.imshow('YOLO Object Detection', image)
+                # Press Q on keyboard to exit
+            if cv2.waitKey(25) & 0xFF == ord('q'):
+                break
+
+        # Break the loop
+        else:
+            break
+    # When everything done, release
+    # the video capture object
+    cap.release()
+    
+    # Closes all the frames
+    cv2.destroyAllWindows()
+    return output_path+'output.avi'
+
+    
+
+if __name__ == '__main__':
+    print(config.config)
+    pipeline("input/6_A_FT_M.mov","output/")