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 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 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 remove_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
# If library directory doesn't exist there's nothing to do
if not os.path.isdir(ACTUAL_DIR):
return
# Remove symbolic links
src_paths, dst_paths = get_src_and_dst_paths(ARDUINO_DIR, subdir)
for dst in dst_paths:
if os.path.islink(dst):
logger.debug("Removing symbolic link {}".format(dst))
os.unlink(dst)
logger.success("Done! :)")
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 save_data_to_file(fields_to_save, out_base_folder, experiment_t_start, filename_prefix, imu_t_start, F_samp):
if F_samp <= 0: # F_samp=0 means data was never collected -> Don't save an empty file!
logger.warn("Nothing to save, skipping saving data to file")
return
out_base_folder = os.path.join(out_base_folder, str(experiment_t_start)[:-7].replace(':', '-'))
os.makedirs(out_base_folder, exist_ok=True) # Make sure base folder exists
out_filename = os.path.join(out_base_folder, "{}{}.h5".format(filename_prefix, str(experiment_t_start)[:-7].replace(':', '-')))
logger.info("Saving collected data at '{}'...".format(out_filename))
with h5py.File(out_filename) as hf:
hf.attrs['t_start'] = str(imu_t_start)
hf.attrs['F_samp'] = F_samp
for field_name, field_children in fields_to_save.items():
field_group = hf.create_group(field_name)
for axis, axis_values in field_children.items():
field_group.create_dataset(axis, data=np.array(axis_values, dtype=np.uint8 if isinstance(axis_values[0], int) else np.float32))
logger.success("Done! Collected data saved at '{}' :)".format(out_filename))
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 main():
T_START = datetime.now()
IMU_SERIAL_PORTS = glob.glob('/dev/cu.SLAB_USBtoUART*') # ['/dev/cu.SLAB_USBtoUART', '/dev/cu.SLAB_USBtoUART9']
imu_ready_queue = Queue(len(IMU_SERIAL_PORTS)) # Use a queue (multi-process safe) to communicate when each process is ready to collect data (esp32 was restarted)
# Start one process per IMU (each BNO055 is connected to its own esp32)
p = []
for i,port in enumerate(IMU_SERIAL_PORTS):
p.append(Process(target=collect_BNO055_data, kwargs={"experiment_t_start": T_START, "serial_port": port, "out_filename_prefix": 'BNO055_{}_'.format(i+1), "imu_ready_queue": imu_ready_queue}))
p[-1].start()
# BLOCK until all esp32s are rebooted and ready to collect data
for i in range(len(p)):
imu_t_start = imu_ready_queue.get()
# Start recording video
record_cam(1, t_start=T_START)
# When record_cam exits (KeyboardInterrupt), terminate all IMU processes and exit
for pp in p:
pp.terminate() # Send a KeyboardInterrupt to each process running collect_BNO055_data
for pp in p:
pp.join() # Wait for all processes to terminate
logger.success("Experiment done, bye!!")
def opened(self):
logger.success("Successfully connected to '{}'!".format(self.url))
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!")
