Darkflow Python Application Example

darkflow-app-example.patch

diff --git a/darkflow/net/flow.py b/darkflow/net/flow.py
index a5b8ceb..bb60704 100644
--- a/darkflow/net/flow.py
+++ b/darkflow/net/flow.py
@@ -90,7 +90,8 @@ def return_predict(self, im):
                 "y": tmpBox[2]},
             "bottomright": {
                 "x": tmpBox[1],
-                "y": tmpBox[3]}
+                "y": tmpBox[3]},
+            "coloridx": tmpBox[5]
         })
     return boxesInfo
 

darkflow-app-example.py

import cv2

from darkflow.net.build import TFNet

Options = {
    'model': 'cfg/tiny-yolo.cfg',
    'load': 'bin/tiny-yolo.weights',
    #"threshold": 0.1
}
DarkflowNet = TFNet(Options)


def drawBoundingBoxes(imageData, imageOutputPath, inferenceResults, colorMap):
    """Draw bounding boxes on an image.

    imageData: image data in numpy array format
    imageOutputPath: output image file path
    inferenceResults: Darkflow inference results
    colorMap: Bounding box color candidates, list of RGB tuples.
    """
    # TODO: return raw data instead of save image
    for res in inferenceResults:
        left = res['topleft']['x']
        top = res['topleft']['y']
        right = res['bottomright']['x']
        bottom = res['bottomright']['y']
        colorIndex = res['coloridx']
        color = colorMap[colorIndex]
        label = res['label']
        confidence = res['confidence']
        imgHeight, imgWidth, _ = imageData.shape
        thick = int((imgHeight + imgWidth) // 300)

        cv2.rectangle(imageData,(left, top), (right, bottom), color, thick)
        cv2.putText(imageData, label, (left, top - 12), 0, 1e-3 * imgHeight,
            color, thick//3)
    cv2.imwrite(imageOutputPath, imageData)


def main():
    imgcv = cv2.imread('./test/horses.jpg')
    results = DarkflowNet.return_predict(imgcv)
    print(results)
    drawBoundingBoxes(imgcv, '/tmp/output.jpg', results, DarkflowNet.meta['colors'])


if __name__ == '__main__':
    main()