def run(self):
# Run forever until event_stop tells us to stop
while not self.data_collection.event_stop.is_set():
# Initialize the websocket
WebSocketBaseClient.__init__(self, self.url, *self.init_args,
**self.init_kwargs)
self.sock.settimeout(
self.TIMEOUT
) # Set the socket timeout so if a host is unreachable it doesn't take 60s (default) to figure out
logger.notice("Connecting to '{}'...".format(self.url))
try:
self.connect() # Attempt to connect to the Arduino
except Exception as e:
logger.error(
"Unable to connect to '{}' (probably timed out). Reason: {}"
.format(self.url, e))
else: # If we were able to connect, then run the websocket (received_message will get called appropriately)
while self.once():
pass # self.once() will return False on error/close -> Only stop when the connection is lost or self.close() is called
self.terminate()
time.sleep(
2
) # Wait for a couple of seconds for Arduino to reboot/connect or just to avoid network overload
logger.success("Thread in charge of '{}' exited :)".format(self.url))
def install_plugins(ARDUINO_DIR):
logger.notice("Installing the SPIFFS (filesystem) plugin...")
call(["./make.sh"],
env=dict(os.environ, INSTALLDIR=ARDUINO_DIR),
cwd="arduino-submodules/plugins/esp32fs")
logger.notice(
"Hopefully everything went right [check the output above] ;)")
def download_hardware_tools():
# Download esp32 tools (not part of the repo so `git submodule update` won't fetch them)
logger.notice(
"Downloading additional esp32 tools not included in the repo...")
try:
esp_tools_folder = os.path.join(ESP_BOARD_FOLDER, "tools")
cur_dir = os.path.abspath(
os.curdir
) # Save a copy of the current dir so we can return after the script is run
os.chdir(esp_tools_folder) # Enter the directory
sys.path.insert(
0,
esp_tools_folder) # Add it to the PATH so get.py can be imported
import get as esp_get # Import the script as a module
## TEMPORARY HACK (hopefully they create a main() function so we don't have to manually copy the contents of the __main__ block)
# Perform the same steps as the "if __name__ == '__main__':" block
esp_get.mkdir_p(esp_get.dist_dir)
tools_to_download = esp_get.load_tools_list(
"../package/package_esp32_index.template.json",
esp_get.identify_platform())
for tool in tools_to_download:
esp_get.get_tool(tool)
# Finally, remember to return to the original folder (so the rest of our script works)
os.chdir(cur_dir)
except Exception as e:
logger.critical("ERROR downloading esp32 tools. Reason: {}".format(e))
return False
logger.success("Additional esp32 tools downloaded :)")
return True
def create_symlinks(ARDUINO_DIR, subdir):
ACTUAL_DIR = os.path.join(ARDUINO_DIR, subdir) # Make it generic so we can symlink Arduino/libraries as well as Arduino/hardware
# Create directory if it doesn't exist
if not os.path.isdir(ACTUAL_DIR):
logger.debug("Arduino '{}' directory does not exist - Creating".format(subdir))
os.makedirs(ACTUAL_DIR)
# Update all libraries (using git submodule)
logger.notice("Making sure you have the latest version of each submodule/library...")
call(["git", "submodule", "init"])
call(["git", "submodule", "update"])
logger.success("All submodules updated :)")
# Create symbolic links
src_paths, dst_paths = get_src_and_dst_paths(ARDUINO_DIR, subdir)
for src, dst in zip(src_paths, dst_paths):
if os.path.exists(dst): # If dst library folder already exists, decide between:
if not os.path.islink(dst): # If the folder is not a symlink and already existed, leave it as is
logger.warning("{} exists and is not a symbolic link - not overwriting".format(dst))
continue
else: # If it was a symlink, just "refresh" (update) it
logger.verbose("Unlinking {} first".format(dst))
os.unlink(dst)
# Create symbolic link
logger.debug("Creating new symbolic link {}".format(dst))
os.symlink(src, dst)
logger.success("Done! :)")
return True
def stop(self):
logger.notice("Stopping data collection!")
self.event_stop.set() # Let the threads know they need to exit
# First, close the websockets
for ws in self.ws_threads:
Thread(target=ws.close).start() # ws.close is blocking so just call it from a new thread (as long as we're not collecting data from too many nodes, we shouldn't hit the max thread limit)
# And wait for all threads to finish
for ws in self.ws_threads:
ws.join()
def record_cam(cam, out_folder='data', t_start=None, save_as_video=True):
if t_start is None: t_start = datetime.now() # Initialize t_start to current time if t_start wasn't specified
if isinstance(t_start, datetime): t_start = str(t_start)[:-7].replace(':', '-') # Convert to str
out_folder = os.path.join(out_folder, t_start)
os.makedirs(out_folder, exist_ok=True) # Ensure folder exists
# Open the camera
video_capture = cv2.VideoCapture(cam)
fps = video_capture.get(cv2.CAP_PROP_FPS)
ret, frame = video_capture.read() # Read a frame, sometimes first read() returns an "invalid" image
if not ret:
logger.critical("Oops, can't access cam {}, exiting! :(".format(cam))
return
frame_dimensions = frame.shape[1::-1] # width, height
# Create a video file if save_as_video
if save_as_video:
video_filename = os.path.join(out_folder, "cam_{}_{}.mp4".format(cam, t_start))
video_writer = cv2.VideoWriter(video_filename, cv2.VideoWriter_fourcc(*'avc1'), fps if fps > 0 else 25, frame_dimensions) # Note: avc1 is Apple's version of the MPEG4 part 10/H.264 standard apparently
logger.notice("Starting cam '{}' recording! Saving {}".format(cam, "as {}".format(video_filename) if save_as_video else "at {}".format(out_folder)))
t_frames = []
try:
while True:
ret, frame = video_capture.read()
t_frames.append(datetime.now())
if len(t_frames) > 1: logger.debug("Wrote frame {} with delta={:.3f}ms (saving took {:.3f}ms)".format(len(t_frames), 1000*(t_frames[-1]-t_frames[-2]).total_seconds(), 1000*(t_save-t_frames[-2]).total_seconds()))
if not ret:
logger.critical("Unknown error capturing a frame from cam {}!")
elif save_as_video:
video_writer.write(frame)
else: # Save as still image
cv2.imwrite(os.path.join(out_folder, "cam_{}_f{:05d}_{}.jpg".format(cam, len(t_frames), t_frames[-1].strftime("%H-%M-%S-%f"))), frame)
t_save = datetime.now()
except KeyboardInterrupt:
logger.notice("Stopping cam recording!")
finally:
video_capture.release()
info_filename = os.path.join(out_folder, "cam_{}_{}.h5".format(cam, t_start))
with h5py.File(info_filename) as hf:
hf.attrs["t_start"] = t_start
hf.attrs["fps"] = fps
hf.attrs["width"] = frame_dimensions[0]
hf.attrs["height"] = frame_dimensions[1]
hf.create_dataset("t_frames", data=np.array([t.timestamp() for t in t_frames]))
logger.success("Goodbye from cam {} recording!".format(cam))
def create_symlinks(ARDUINO_DIR, subdir):
ACTUAL_DIR = os.path.join(ARDUINO_DIR, subdir) # Make it generic so we can symlink Arduino/libraries as well as Arduino/hardware
# Create directory if it doesn't exist
if not os.path.isdir(ACTUAL_DIR):
logger.debug("Arduino '{}' directory does not exist - Creating".format(subdir))
os.makedirs(ACTUAL_DIR)
# Update all libraries (using git submodule)
logger.notice("Making sure you have the latest version of each submodule/library...")
call(["git", "submodule", "init"])
call(["git", "submodule", "update"])
logger.success("All submodules updated :)")
# Create symbolic links
src_paths, dst_paths = get_src_and_dst_paths(ARDUINO_DIR, subdir)
for src, dst in zip(src_paths, dst_paths):
make_symlink(src, dst)
logger.success("Done! :)")
return True
def received_message(self, msg):
if msg.is_text:
return # Ignore Geophone ID message (eg: Geophone_AABBBCC)
# Parse the message: '<' for Little-Endian, 'H' for uint16_t
msg_format = '<' + 'H' * (len(msg.data) / 2)
msg_vals = unpack(msg_format, msg.data)
msg_id = msg_vals[0]
msg_vals = msg_vals[1:]
cvs_vals = ','.join(map(
str, msg_vals)) # Convert each item to str then join with ','
# Check if we need to start a new file
if datetime.now() > self.deadline_new_file:
# Close existing file if necessary
if self.output_file_handle:
self.output_file_handle.close()
logger.verbose("Closed file: '{}' (it's been {}s)".format(
self.output_filename, self.DELTA_NEW_FILE.total_seconds()))
# And create a new one
self.generate_new_filename()
self.output_file_handle = open(self.output_filename, 'w')
# Write the parsed message to the file
try: # In case the file has been closed (user stopped data collection), surround by try-except
self.output_file_handle.write(cvs_vals + ',')
except Exception as e:
logger.error("Couldn't write to '{}'. Error: {}".format(
self.output_filename, e))
if self.last_msg_id is not None and msg_id != self.last_msg_id + 1:
logger.notice(
"Received data from '{}' (msg ID {})! WARNING: {} packet(s) lost: last msg ID was {}!"
.format(self.url, msg_id, msg_id - self.last_msg_id - 1,
self.last_msg_id))
else:
logger.debug("Received data from '{}' (msg ID {})!".format(
self.url, msg_id))
self.last_msg_id = msg_id
def collect_BNO055_data(experiment_t_start=None, serial_port='/dev/cu.SLAB_USBtoUART', baud_rate=115200, out_base_folder='data', out_filename_prefix='BNO055_', imu_ready_queue=None):
# NOTE: In order to allow for multiple simultaneous IMUs, experiment_t_start indicates the name of the enclosing folder where the data will be saved,
# and each one will have their own imu_t_start indicating the precise time when each IMU was rebooted (and therefore its data started logging)
logger.info("Connecting to {}...".format(serial_port))
s = serial.Serial(serial_port, baud_rate)
logger.success("Successfully connected! Please reboot the ESP32")
# Initialize variables
fields_to_save = init_fields_to_save()
data_block = SensorData.DataBlock()
imu_t_start = None
F_samp = 0
try:
# Read serial until we receive the " -- SETUP COMPLETE -- " message
while True:
l = s.readline().rstrip() # Read a line (removing any trailing \r\n)
logger.debug(l)
if l == b'-- SETUP COMPLETE --':
logger.notice("Esp32 is done booting, collecting data until Ctrl+C is pressed!")
break
# Update imu_t_start to the actual time the data collection started
imu_t_start = datetime.now()
if experiment_t_start is None:
experiment_t_start = imu_t_start # If user didn't specify an experiment t_start, it's probably only collecting data for 1 IMU -> Use imu_t_start to name the data folder
if imu_ready_queue is not None: # Pass the current time through the queue to notify the main process to start recording the camera
imu_ready_queue.put(imu_t_start)
# Collect data (read samples and append them to fields_to_save)
last_msg_id = None
while True:
pb_len = struct.unpack('<H', s.read(2))[0] # Read 2 bytes that indicate how long the proto message will be and convert to uint16
pb = s.read(pb_len) # Read the proto contents
data_block.ParseFromString(pb) # Decode them
logger.info("Message received: {} (t={:.6f})".format(data_block.id, data_block.t_latest.seconds + data_block.t_latest.nanos/1e9))
# Check if any data got lost (hopefully not :D)
if last_msg_id is None:
F_samp = data_block.F_samp # Store the sampling frequency so we can save it in the file
elif data_block.id - last_msg_id > 1:
n_lost = data_block.id - last_msg_id - 1
logger.critical("Lost {} packet{}! (Previous message ID was {}, current is {})".format(n_lost, 's' if n_lost>1 else '', data_block.id, last_msg_id))
last_msg_id = data_block.id
# Append new data to fields_to_save
for field_name, field_children in fields_to_save.items():
field_children_values = getattr(data_block, field_name)
for axis, axis_values in field_children.items():
axis_new_values = getattr(field_children_values, axis) # Fetch new sensor readings
axis_values.append(axis_new_values) if isinstance(axis_new_values, int) else axis_values.extend(axis_new_values) # Append values
except KeyboardInterrupt:
logger.notice("Stopping data collection!")
finally:
# Close serial port and save all data to a file
s.close()
save_data_to_file(fields_to_save, out_base_folder, experiment_t_start, out_filename_prefix, imu_t_start, F_samp)
logger.success("Goodbye from BNO055 data collection!")
def build_protos():
logger.notice("Building message protobufs...")
call(["./build-protos.sh"], cwd="protos")
logger.notice(
"Hopefully everything went right [check the output above] ;)")