forked from singnet/dnn-model-services
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_service.py
65 lines (55 loc) · 2.31 KB
/
test_service.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
import sys
import grpc
# import the generated classes
import service.service_spec.scene_recognition_pb2_grpc as grpc_bt_grpc
import service.service_spec.scene_recognition_pb2 as grpc_bt_pb2
from service import registry
import json
from service.serviceUtils import base64_to_jpg
TEST_URL = "https://raw.githubusercontent.com/singnet/dnn-model-services/master/docs/assets/users_guide/places365_beach.png"
if __name__ == "__main__":
try:
test_flag = False
if len(sys.argv) == 2:
if sys.argv[1] == "auto":
test_flag = True
# If running a test script, fill request automatically
if test_flag:
# Open a gRPC channel to endpoint given by registry
endpoint = "localhost:{}".format(registry["scene_recognition_service"]["grpc"])
# Fill request
grpc_method = "recognize_scene"
input_image = TEST_URL
predict = "IO, Attributes, Categories, CAM"
else:
endpoint = input("Endpoint")
grpc_method = "recognize_scene"
input_image = input("Image URL: ")
predict = input("Pick what to predict (csv: io, attributes, categories, cam): ")
if grpc_method == "recognize_scene":
# Open a channel to grpc endpoint
channel = grpc.insecure_channel(endpoint)
print("Opened channel")
# Create a stub (client)
stub = grpc_bt_grpc.SceneRecognitionStub(channel)
# Create a valid request message
request = grpc_bt_pb2.SceneRecognitionRequest(input_image=input_image,
predict=predict)
# Make the call
response = stub.recognize_scene(request)
# Print and treat response
response_dict = json.loads(response.data)
print("TEST RESULT: ")
for key, value in response_dict.items():
if key == "cam":
base64_to_jpg(value, "./test_cam.jpg")
print("CAM saved to ./test_cam.jpg")
else:
print("{}: {}".format(key, value))
print("Service completed!")
else:
print("Invalid method!")
exit(1)
except Exception as e:
print(e)
exit(1)