Created
April 22, 2017 10:21
-
-
Save bafu/74737f2c70df17609833a6a9ab35978a to your computer and use it in GitHub Desktop.
Darkflow Python Application Example
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment