1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
|
# Copyright 2017 DT42
#
# This file is part of BerryNet.
#
# BerryNet is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# BerryNet is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with BerryNet. If not, see <http://www.gnu.org/licenses/>.
"""Darkflow inference engine.
"""
from __future__ import print_function
import argparse
import logging
import cv2
from darkflow.net.build import TFNet
from engineservice import EngineService
from engineservice import DLEngine
from dlmodelmgr import DLModelManager
# FIXME: Make these variables configurable
SystemSnapshot = '/usr/local/berrynet/dashboard/www/freeboard/snapshot.jpg'
class DarkflowEngine(DLEngine):
def __init__(self, model, label, config):
super(DarkflowEngine, self).__init__()
self.engine_options = {
'model': config,
'load': model,
#'model': "cfg/tiny-yolo.cfg",
#'load': "bin/tiny-yolo.weights",
'verbalise': True,
"threshold": 0.1
}
def create(self):
self.tfnet = TFNet(self.engine_options)
def inference(self, tensor):
return self.tfnet.return_predict(tensor)
def save_cache(self):
#with open(self.cache['model_output_filepath'], 'w') as f:
# f.write(str(self.cache['model_output']))
drawBoundingBoxes(self.cache['model_input'],
#self.cache['model_output_filepath'] + '.jpg',
SystemSnapshot,
self.cache['model_output'],
self.tfnet.meta['colors'])
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)
logging.debug('Save bounding box result image to {}'.format(imageOutputPath))
def parse_args():
ap = argparse.ArgumentParser()
ap.add_argument('--model',
help='Model file path')
ap.add_argument('--label',
help='Label file path')
ap.add_argument('--model_package',
default='',
help='Model package "name-version" naming')
ap.add_argument('--image_dir', required=True,
help='Path to image file')
ap.add_argument('--service_name', required=True,
help='Engine service name used as PID filename')
ap.add_argument('--num_top_predictions', default=5,
help='Display this many predictions')
return vars(ap.parse_args())
if __name__ == '__main__':
logging.basicConfig(level=logging.DEBUG)
args = parse_args()
if args['model_package'] != '':
dlmm = DLModelManager()
meta = dlmm.get_model_meta(args['model_package'])
args['model'] = meta['model']
args['label'] = meta['label']
args['config'] = meta['config']['graph']
logging.debug('model filepath: ' + args['model'])
logging.debug('label filepath: ' + args['label'])
logging.debug('image_dir: ' + args['image_dir'])
darkflow_engine = DarkflowEngine(args['model'], args['label'], args['config'])
engine_service = EngineService(args['service_name'], darkflow_engine)
engine_service.run(args)
# this code block works
#import cv2
#input_tensor = cv2.imread('/tmp/berrynet/dog.jpg')
#tensor = darkflow_engine.process_input(input_tensor)
#output = darkflow_engine.inference(tensor)
#output = darkflow_engine.process_output(output)
|
