-
Notifications
You must be signed in to change notification settings - Fork 0
/
client.py
185 lines (169 loc) · 9.36 KB
/
client.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
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
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
import argparse
from datetime import datetime
import locale
import os
from pathlib import Path
import pickle
import socket
import sys
import json
# CLI Arguments
parser = argparse.ArgumentParser(description="Client for sending images to an object recogniction service")
parser.add_argument("-f", "--file", help="Path to file, also accepts pipeline input", default=None) # accepts pipeline input
parser.add_argument("-s", "--server", help="The host/IP address of the server, required.", required=True)
parser.add_argument("-p", "--port", help="TCP Port to send to. | Example (default): 4949", type=int, default=4949)
parser.add_argument("-b", "--buffer", help="Buffer size. | Example (default): 4096", type=int, default=4096)
parser.add_argument("-e", "--encoding", help="Set encoding. | Example (default): utf-8", default='utf-8')
parser.add_argument("-c", "--objectclass", help="Object class(s) to recognize. | Example 1: cat | Example 2: car airplane 'potted plant' chair | Notes: When using multiple object classes, --maxresults is ignored", nargs='+', default=None)
parser.add_argument("-x", "--detectionbox", help="Detection box size (0.0 min - 1.0 max) [y_min, x_min, y_max, x_max] | Example : 0.0 0.0 0.5 0.5", nargs='+', default=None)
parser.add_argument("-m", "--minscore", help="Minimum detection score (percent). | Example (default): 60", type=int, default=60)
parser.add_argument("-a", "--maxresults", help="Max results returned. | Example (default): 1", type=int, default=1)
parser.add_argument("-u", "--output", help="Output type: silent, simple, detailed, json, onlyclass, onlyscore, onlybox, onlyclassfound, onlywithinbox, raw. | Example (default): simple", default="simple")
parser.add_argument("-d", "--logging", help="Output logging information.", default=False, action='store_true')
args = parser.parse_args()
# Locale
locale.setlocale(locale.LC_ALL, '')
if args.logging: print(f'{datetime.now().strftime("%c")} | {args}')
# Socket
SERVER_HOST = args.server # Server IP Address
SERVER_PORT = args.port # Server Port (tcp)
BUFFER_SIZE = args.buffer # Bytes to receive each time
SEPARATOR = "###"
# Client Parameters
if isinstance(args.objectclass, list):
fileObjectClass = list((map(lambda x: x.lower(), args.objectclass)))
fileMaxResults = len(fileObjectClass)
elif args.objectclass == None:
fileObjectClass = None
fileMaxResults = args.maxresults
fileBox = args.detectionbox
fileMinScore = args.minscore
# Functions
def sendFile(fileName):
# Create socket and connect to host and port
clientSocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
clientSocket.connect((SERVER_HOST, SERVER_PORT))
clientSocket.recv(BUFFER_SIZE).decode(args.encoding)
if args.logging: print(f'{datetime.now().strftime("%c")} | {clientSocket.getsockname()}')
# Send file info
fileSize = os.path.getsize(fileName)
baseName = os.path.basename(fileName)
fileSizeLocaleRound = locale.format_string('%d', round(fileSize / 1024 , 1), grouping=True)
if args.logging: print(f'{datetime.now().strftime("%c")} | Sending File: {fileName}')
if args.logging: print(f'{datetime.now().strftime("%c")} | Size: {fileSizeLocaleRound} KB')
if args.logging: print(f'{datetime.now().strftime("%c")} | Uploading...')
clientSocket.send(f"{fileName}{SEPARATOR}{fileSize}".encode(args.encoding).strip()) # Send file info (name and size)
clientSocket.recv(BUFFER_SIZE).decode(args.encoding)
# Send file
with open(fileName, "rb") as f:
while True:
bytesRead = f.read(BUFFER_SIZE) # Read the bytes from the file
if not bytesRead:
break # Done sending file
clientSocket.sendall(bytesRead)
# Receive data back
modelInfo = clientSocket.recv(BUFFER_SIZE).decode(args.encoding)
clientSocket.send("ACK".encode(args.encoding))
if args.logging: print(f'{datetime.now().strftime("%c")} | TensorFlow Model: {modelInfo}')
pickledInferenceResults = clientSocket.recv(BUFFER_SIZE)
clientSocket.send("ACK".encode(args.encoding))
inferenceResults = pickle.loads(pickledInferenceResults)
pickledCategoryIndex = clientSocket.recv(BUFFER_SIZE)
clientSocket.send("ACK".encode(args.encoding))
categoryIndex = pickle.loads(pickledCategoryIndex)
# Store results
resultNumDetections, allResultClasses, allResultScores, allResultBoxes = inferenceResults['num_detections'][0].astype(int), inferenceResults['detection_classes'][0], inferenceResults["detection_scores"][0], inferenceResults["detection_boxes"][0]
# Begin outputting
silentExit = 1
if args.output == "raw":
print(inferenceResults)
else:
jsonOutput = []
for i in range(resultNumDetections):
if allResultScores[i] >= fileMinScore/100 and i < fileMaxResults:
# Determine if object in detection box
if fileBox != None: # [y_min, x_min, y_max, x_max]
boxCenterX = (allResultBoxes[i][1]+allResultBoxes[i][3])/2
boxCenterY = (allResultBoxes[i][0]+allResultBoxes[i][2])/2
if (boxCenterX >= float(fileBox[1])) and (boxCenterX <= float(fileBox[3])) and (boxCenterY >= float(fileBox[0])) and (boxCenterY <= float(fileBox[2])):
withinBox = True
else:
withinBox = False
# Only output when objects are within box
if args.output != "silent":
continue
else:
withinBox = None
# Determine if object class found
if fileObjectClass != None:
for currentObjectClass in range(len(fileObjectClass)):
if fileObjectClass[currentObjectClass] == categoryIndex[allResultClasses[i]]["name"]:
resultClassFound = True
break
else:
resultClassFound = False
# Only output when classes that were specified are found
if resultClassFound == False and args.output != "silent":
continue
else:
resultClassFound = None
# Print output
if args.output == "silent":
if resultClassFound == True and withinBox == None: # Class only
silentExit = 0
break
elif resultClassFound == None and withinBox == True: # Box only
silentExit = 0
break
elif resultClassFound == True and withinBox == True: # Class & Box
silentExit = 0
break
else:
silentExit = 1
elif args.output == "simple":
print(f'{baseName}, {categoryIndex[allResultClasses[i]]["name"]}, {int(allResultScores[i] * 100)}, {allResultBoxes[i]}, {resultClassFound}, {withinBox}')
elif args.output == "detailed":
print(f'Filename: {baseName}')
print(f'Inferenced Class: {categoryIndex[allResultClasses[i]]["name"]}')
print(f'Inferenced Score: {int(allResultScores[i] * 100)}%')
print(f'Inferenced Box: {allResultBoxes[i]}')
print(f'Class Found: {resultClassFound}')
print(f'Within Detection Box: {withinBox}')
elif args.output == "json":
jsonOutput.append({ "Filename":baseName, "Inferenced Class":categoryIndex[allResultClasses[i]]["name"], "Inferenced Score":int(allResultScores[i] * 100), "Inferenced Box":str(allResultBoxes[i]), "Class Found":resultClassFound, "Within Detection Box":withinBox})
elif args.output == "onlyclass":
print(f'{categoryIndex[allResultClasses[i]]["name"]}')
elif args.output == "onlyscore":
print(f'{int(allResultScores[i] * 100)}%')
elif args.output == "onlybox":
print(f'{allResultBoxes[i]}')
elif args.output == "onlyclassfound":
print(f'{resultClassFound}')
elif args.output == "onlywithinbox":
print(f'{withinBox}')
else:
print(f'Unknown output option specified.')
# Print JSON output
if args.output == "json":
print(json.dumps(jsonOutput, indent=3))
# Close the client socket
if args.logging: print(f'{datetime.now().strftime("%c")} | Closing Connection')
clientSocket.close()
# Check for silent output
if args.output == "silent":
#print(silentExit)
sys.exit(silentExit)
# Check for pipeline input, folder, or file and call sendFile()
if args.file == None:
for fileName in sys.stdin:
sendFile(fileName.strip())
elif os.path.isdir(args.file):
folderName = Path(args.file)
allFiles = list(folderName.glob("*"))
for fileName in allFiles:
sendFile(fileName)
elif os.path.isfile(args.file):
fileName = args.file
sendFile(fileName)
else:
print("Unknown file input.")