-
Notifications
You must be signed in to change notification settings - Fork 0
/
server.py
executable file
·107 lines (79 loc) · 3.23 KB
/
server.py
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
import argparse
import common
import cv2
import numpy as np
import time
import logging
from gabriel_protocol import gabriel_pb2
from gabriel_server import local_engine
from gabriel_server import cognitive_engine
logging.basicConfig(filename="Loglatency.log", level=logging.INFO)
net = cv2.dnn.readNet('yolov3.weights','yolov3.cfg')
classes = []
with open('coco.names','r') as f:
classes = f.read().splitlines()
font = cv2.FONT_HERSHEY_PLAIN
colors = np.random.uniform(0, 255, size=(100, 3))
def getresults(img_str):
t = time.time()
nparr = np.fromstring(img_str, np.uint8)
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) # cv2.IMREAD_COLOR
height, width,_ = img.shape
blob = cv2.dnn.blobFromImage(img, 1 / 255, (416, 416), (0, 0, 0), swapRB=True, crop=False)
net.setInput(blob)
output_layers_names = net.getUnconnectedOutLayersNames()
layerOutputs = net.forward(output_layers_names)
boxes = []
confidences = []
class_ids = []
for output in layerOutputs:
for detection in output:
scores = detection[5:]
class_id = np.argmax(scores)
confidence = scores[class_id]
if confidence > 0.2:
center_x = int(detection[0] * width)
center_y = int(detection[1] * height)
w = int(detection[2] * width)
h = int(detection[3] * height)
x = int(center_x - w / 2)
y = int(center_y - h / 2)
boxes.append([x, y, w, h])
confidences.append((float(confidence)))
class_ids.append(class_id)
indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.2, 0.4)
if len(indexes) > 0:
for i in indexes.flatten():
x, y, w, h = boxes[i]
label = str(classes[class_ids[i]])
confidence = str(round(confidences[i], 2))
color = colors[i]
cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)
cv2.putText(img, label + " " + confidence, (x, y + 20), font, 2, (255, 255, 255), 2)
t1 = time.time()
tf = t1 - t
logging.info('{}'.format(tf))
return cv2.imencode('.jpg', img)[1].tostring()
class DisplayEngine(cognitive_engine.Engine):
def handle(self, input_frame):
status = gabriel_pb2.ResultWrapper.Status.SUCCESS
result_wrapper = cognitive_engine.create_result_wrapper(status)
image = getresults(input_frame.payloads[0])
result = gabriel_pb2.ResultWrapper.Result()
result.payload_type = gabriel_pb2.PayloadType.IMAGE
#result.payload = input_frame.payloads[0]
result.payload = image
result_wrapper.results.append(result)
return result_wrapper
def main():
common.configure_logging()
parser = argparse.ArgumentParser()
parser.add_argument(
'source_name', nargs='?', default=common.DEFAULT_SOURCE_NAME)
args = parser.parse_args()
def engine_factory():
return DisplayEngine()
local_engine.run(engine_factory, args.source_name, input_queue_maxsize=60,
port=common.WEBSOCKET_PORT, num_tokens=2)
if __name__ == '__main__':
main()