-
Notifications
You must be signed in to change notification settings - Fork 0
/
unitycommunicator.py
428 lines (325 loc) · 15.1 KB
/
unitycommunicator.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
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
import socket
import json
import logging
import numpy as np
from PIL import Image
import io
import subprocess
class UnityCommunicator:
"""
Provides all methods for communication to Unity 3D game enginge.
Scene parameters can be sent to unity, the corresponding screen shot of the scene will be sent back. Therefore a tcp connection is used.
The tcp parameters such as 'host' and 'port' are specified in tcpconfig.json file in the unity projects 'StreamingAssets'
folder. This class is best used in a with-statement.
See also
--------
Documentation 'unitycommunicator'
Created during software practical at Heidelberg Collaboratory of image processing. September 2019
https://hciweb.iwr.uni-heidelberg.de/compvis
Examples
--------
Render a scene parameter set (from json dictionary which can be passed directly or read from file).
>>> from unitycommunicator import UnityCommunicator
>>> with UnityCommunicator('unity_build.app', use_with_unity_build=True, log_level=logging.WARNING, width=600, height=400) as uc:
>>> json_data = uc.read_json_file('parameterfile.json')
>>> scene_img, scene_id = uc.render_parameters(json_data)
To save image, use e.g. pillow
>>> from PIL import Image
>>> Image.fromarray(scene_img).save('Rendered_Scene_ID-{:3}.png'.format(str(scene_id).zfill(3)))
>>> # str.zfill ensures a consistent file naming (leading zeros)
Dependencies
------------
Libraries, that are used by this class.
socket
json
logging
numpy
PIL.Image
io
os
"""
def __init__(self, unity_build_path: str, use_with_unity_build=True, log_level=logging.WARNING, width=600, height=400):
"""
Sets up everything to enable a tcp connection to unity.
Constructor of class, starts logging, starts unity build executable, starts tcp socket, sends resolution for rendering.
Parameters
---------
unity_build_path : str
path to unity build executable
use_with_unity_build : bool
if used with unity engine True, if used with unity build False
Changes some path specifications and also starts the build if True.
Default: True
log_level : logging.<LEVEL>
log level of the console output.
Default: logging.warning.
width : int
width of rendered picture in px.
height : int
height of rendered picture in px.
"""
log_path = 'log/unity-communicator.log'
# Create logger
self.logger = logging.getLogger('UClog')
self.logger.setLevel(logging.DEBUG)
# Create console handler
self.ch = logging.StreamHandler()
self.ch.setLevel(log_level)
# Create file handler
self.fh = logging.FileHandler(log_path)
self.fh.setLevel(logging.DEBUG)
# Add formatter
self.formatter_fh = logging.Formatter('%(asctime)s - %(levelname)s : %(message)s')
self.fh.setFormatter(self.formatter_fh)
self.formatter_ch = logging.Formatter('%(levelname)s : %(message)s')
self.ch.setFormatter(self.formatter_ch)
# Add fh and ch to logger
self.logger.addHandler(self.fh)
self.logger.addHandler(self.ch)
#Clear log at startup if it is longer than 1000 lines
with open(log_path, 'r') as f:
log_length = len(f.readlines())
if log_length > 1000:
with open(log_path, 'w'):
pass
print('')
self.logger.info('starting unity communication server...\n')
self.unity_build_path = unity_build_path
if use_with_unity_build == True:
# For execution with BUILD
# Start unity build
try:
subprocess.call(["open", self.unity_build_path])
except OSError:
self.logger.error('Starting Unity build failed. Please start it manually.')
self.streaming_assets_path = self.unity_build_path + '/Contents/Resources/Data/StreamingAssets/'
else:
# For execution with unity engine
self.streaming_assets_path = self.unity_build_path + '/Assets/StreamingAssets/'
self.logger.warning('Please start Unity...')
# Specify paths to tcpconfig.json file (in streamingAssets folder of unity project
self.tcp_config_path = self.streaming_assets_path + 'tcpconfig.json'
# Timeout for socket connection, will be used to prevent lock (in case both peers are receiving)
self.socket_timeout = 0.1
# Call function for tcp server setup
self.conn_unity = self._setup_server()
self.conn_closed = False
# Send resolution to unity
unity_init_dict = {'width':width, 'height':height}
unity_init_string = json.dumps(unity_init_dict, separators=(',', ':'))
self._send_data(message=unity_init_string)
#Receive confirmation
unity_response = self._receive_complete_message().decode()
self.logger.warning(unity_response + '\n')
def __enter__(self):
"""
Necessary for usage in with...as statement.
Returns
-------
self : UnityCommunicator
For use in "with UnityCommunicator(unity_build_path) as:"
"""
return self
def __exit__(self, type, value, traceback):
"""Called when used in with statement. Sends end request to Unity, closes TCP connection. """
self.close()
self.conn_closed = True
def __del__(self):
"""Class destructor. Sends end request to Unity, closes TCP connection. """
if not self.conn_closed:
self.close()
self.conn_closed = True
def close(self):
"""sends end request to Unity, closes TCP connection."""
if not self.conn_closed:
self._send_data('END.')
self.logger.warning('End command was sent to Unity.')
self.conn_unity.close()
self.logger.warning('Socket closed. Shutting down...')
self.conn_closed = True
def _setup_server(self):
"""
Sets up TCP server for unity connection.
Therefore it looks for a free port in range (50000, 50099) and stores
the used port in tcpconfig.json file of Unity project at /StreamingAssets/tcpconfig.json.
Unity reads the port and accepts the connection.
Returns
--------
socket.socket
the socket connection to unity
Raises
------
FileNotFoundError
tcpconfig.json in StreamingAssets Folder of UnityProject can not be found
socket.error
to catch if socket is in use
"""
# Open tcpconfig.json file
try:
with open(self.tcp_config_path, 'r') as f:
config = json.load(f)
f.close()
except FileNotFoundError as e:
self.logger.error('\'tcpconfig.json\' no such file, should be found in <UnityProject>/Assets/StreamingAssets/')
self.conn_closed = True # To prevent __del__() from closing a not existing connection
raise e
# try to bind port specified in tcpconfig.json file
while 1:
try:
socket_unity = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
socket_unity.bind((config['host'], config['ports'][1]))
except socket.error as e:
# If port is in use, try next port
print(e)
self.logger.warning('try next port...')
if (config['ports'][1]) >= 50100:
config['ports'][1] = 50000
else:
config['ports'][1] += 1
else:
break
# Write used port to 'tcpconfig.json'
with open(self.tcp_config_path, 'w') as f:
json.dump(config, f)
f.close()
socket_unity = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
socket_unity.bind((config['host'], config['ports'][1]))
socket_unity.listen(1)
self.logger.warning('waiting for connection from Unity at port {} ...'.format(config['ports'][1]))
conn_unity, addr_unity = socket_unity.accept()
self.logger.warning('connection to Unity established at addr: ' + str(addr_unity))
conn_unity.settimeout(self.socket_timeout)
return conn_unity
def _send_data(self, message: str):
"""
Sends string to socket.
At the end of each message, an end tag 'eod.' is added to the message.
Unity detects the end of the message by reading this tag.
Parameters
---------
message : str
server message to unity. should use json format.
"""
self.conn_unity.sendall((message + 'eod.').encode()) # End of Data
self.logger.debug('_send_data(): sendall succeeded')
def _receive_data_as_bytes(self):
"""
Receives data at the classes socket.
_receive_data_as_bytes is either ended by timeout or if not data is sent.
Returns
-------
data_complete : bytearray
bytearray of message packages
"""
self.logger.debug('_receive_data_as_bytes(): entered')
data_complete = bytearray([0])
while 1:
try:
data = self.conn_unity.recv(1024)
except socket.timeout:
self.logger.debug('_receive_data_as_bytes(): timeout -> exit _receive_data_as_bytes()')
break
data_complete = data_complete + data
if not data:
self.logger.debug('_receive_data_as_bytes(): no data anymore. receive ends.')
break
return data_complete
def _receive_complete_message(self):
"""
Receives data until the valid end tag of the message is detected, then return message as bytes.
Returns
-------
unity_resp_bytes : bytearray
bytearray of a complete message from unity
"""
# Run the _receive_data_as_bytes() until valid response (bytearray end tag).
# This uses the socket timeout break in _receive_data_as_bytes() to wait for the rendering to finish.
# This is necessary to avoid a lock, which occurs when python starts to receive before unity sends.
while 1:
unity_resp_bytes = self._receive_data_as_bytes()
if unity_resp_bytes[-8:] == bytearray([255, 0, 250, 251, 252, 253, 254, 255]):
self.logger.debug('_receive_complete_message(): End tag from Unity detected, end receive')
break
self.logger.debug('_receive_complete_message(): Run again receive_data()')
return unity_resp_bytes[:-8]
def read_json_file(self, file_name: str):
"""
Loads json file as python directory.
Parameters
---------
filename : str
path to parameter file
Returns
-------
dictionary
contains json file parameters. Returns None if invalid file or parsing to dict impossible.
"""
try:
json_string = open(file_name).read()
except FileNotFoundError as e:
self.logger.error(file_name + ' can not be found. returned None')
#Return none, in order to skip this parameter file. (None check is implemented in render_parameters())
return None
try:
json_dict = json.loads(json_string)
except json.decoder.JSONDecodeError as e:
self.logger.error(file_name + ' can not be converted to json. Returned None')
# Return none, in order to skip this parameter file. (None check is implemented in render_parameters())
return None
return json_dict
def render_parameters(self, param_dict: dict):
"""
Controls the rendering of one scene parameter set in unity.
Therefore it delivers given parameters as json string to unity and returns the scene meta data as well as a
rendered scene picture. The json_dict should contain parameters for every scene object in unity. Also the class
'JSONCaptureParameters' in Assets/Scripts/ in unity should contain a representation for the delivered dict.
Parameters
---------
param_dict : dictionary
should contain all parameters of scene objects in unity. Also the unchanged ones compared to the last frame.
Returns
-------
np.array
contains captured scene as RGBA array of dimensions (height, width, 4)
int
the sceneID specified in the deliverd param_dict for the scene
See also
--------
Documentation unity communicator.
Examples
--------
See UnityCommunicator().__doc__
"""
if param_dict is None:
self.logger.error('render_parameters(): Aborted because of invalid json_dict')
return None
self.logger.info('Render parameters of scene with ID ' + str(param_dict['sceneID']) + ' ...')
json_string = json.dumps(param_dict, separators=(',', ':'))
self.logger.debug('render_parameters(): Start sending data...')
self._send_data(message=json_string)
self.logger.debug('render_parameters(): Success! Data sent.')
unity_resp_bytes = self._receive_complete_message()
self.logger.debug('render_parameters(): Received scene of ' + str(len(unity_resp_bytes)) + ' bytes')
# Load the meta data file in dictionary and process the scene_img
meta_length = int.from_bytes(unity_resp_bytes[-4:], byteorder='little')
meta_bytes = unity_resp_bytes[-(4 + meta_length):-4].decode()
meta_data_dict = json.loads(meta_bytes)
scene_id = meta_data_dict['sceneID']
message = meta_data_dict['message']
img_bytes = unity_resp_bytes[1:-(12+meta_length)]
self.logger.info('Received and unpacked scene with ID ' + str(scene_id))
self.logger.info('Scene meta data message: ' + message + '\n')
pillow_img = Image.open(io.BytesIO(img_bytes))
scene_img = np.array(pillow_img)
return scene_img, meta_data_dict
if __name__ == '__main__':
with UnityCommunicator('UCGeometrics', use_with_unity_build=False, log_level=logging.INFO, width=550, height=400) as uc:
json_data = uc.read_json_file('ParameterFiles/parameters_geometrics0.json')
scene_img, meta_data_dict = uc.render_parameters(json_data)
Image.fromarray(scene_img).save('SavedScenes/Rendered_Scene_ID-{:3}.png'.format(str(meta_data_dict['sceneID']).zfill(3)))
json_data = uc.read_json_file('ParameterFiles/parameters_geometrics1.json')
scene_img, meta_data_dict = uc.render_parameters(json_data)
Image.fromarray(scene_img).save('SavedScenes/Rendered_Scene_ID-{:3}.png'.format(str(meta_data_dict['sceneID']).zfill(3)))
json_data = uc.read_json_file('ParameterFiles/parameters_geometrics2.json')
scene_img, meta_data_dict = uc.render_parameters(json_data)
Image.fromarray(scene_img).save('SavedScenes/Rendered_Scene_ID-{:3}.png'.format(str(meta_data_dict['sceneID']).zfill(3)))