wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

shutil.rmtree(tmp_imgs_dir)

This commit is contained in:
Remi Cadene 2024-07-15 14:19:52 +02:00
parent 03d778e672
commit 87e39da641
1 changed files with 20 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
lerobot/scripts/control_robot.py
View File

@ -173,6 +173,14 @@ def log_control_info(robot, dt_s, episode_index=None, frame_index=None):
logging.info(" ".join(log_items)) logging.info(" ".join(log_items))
def get_is_headless():
if platform.system() == "Linux":
display = os.environ.get("DISPLAY")
if display is None or display == "":
return True
return False
######################################################################################## ########################################################################################
# Control modes # Control modes
######################################################################################## ########################################################################################
@ -241,6 +249,8 @@ def record_dataset(
else: else:
episode_index = 0 episode_index = 0
is_headless = get_is_headless()
# Execute a few seconds without recording data, to give times # Execute a few seconds without recording data, to give times
# to the robot devices to connect and start synchronizing. # to the robot devices to connect and start synchronizing.
timestamp = 0 timestamp = 0
@ -255,6 +265,9 @@ def record_dataset(
now = time.perf_counter() now = time.perf_counter()
observation, action = robot.teleop_step(record_data=True) observation, action = robot.teleop_step(record_data=True)
if not is_headless:
image_keys = [key for key in observation if "image" in key]
dt_s = time.perf_counter() - now dt_s = time.perf_counter() - now
busy_wait(1 / fps - dt_s) busy_wait(1 / fps - dt_s)
@ -270,15 +283,8 @@ def record_dataset(
rerecord_episode = False rerecord_episode = False
stop_recording = False stop_recording = False
def is_headless():
if platform.system() == "Linux":
display = os.environ.get("DISPLAY")
if display is None or display == "":
return True
return False
# Only import pynput if not in a headless environment # Only import pynput if not in a headless environment
if is_headless(): if is_headless:
logging.info("Headless environment detected. Keyboard input will not be available.") logging.info("Headless environment detected. Keyboard input will not be available.")
else: else:
from pynput import keyboard from pynput import keyboard
@ -400,9 +406,11 @@ def record_dataset(
with open(rec_info_path, "w") as f: with open(rec_info_path, "w") as f:
json.dump(rec_info, f) json.dump(rec_info, f)
is_last_episode = stop_recording or (episode_index == (num_episodes - 1))
# Wait if necessary # Wait if necessary
with tqdm.tqdm(total=reset_time_s, desc="Waiting") as pbar: with tqdm.tqdm(total=reset_time_s, desc="Waiting") as pbar:
while timestamp < reset_time_s and not stop_recording: while timestamp < reset_time_s and not is_last_episode:
time.sleep(1) time.sleep(1)
timestamp = time.perf_counter() - start_time timestamp = time.perf_counter() - start_time
pbar.update(1) pbar.update(1)
@ -417,11 +425,10 @@ def record_dataset(
episode_index += 1 episode_index += 1
# Only for last episode if is_last_episode:
if stop_recording or episode_index == num_episodes:
logging.info("Done recording") logging.info("Done recording")
os.system('say "Done recording"') os.system('say "Done recording"')
if not is_headless(): if not is_headless:
listener.stop() listener.stop()
logging.info("Waiting for threads writing the images on disk to terminate...") logging.info("Waiting for threads writing the images on disk to terminate...")
@ -446,6 +453,7 @@ def record_dataset(
# note: `encode_video_frames` is a blocking call. Making it asynchronous shouldn't speedup encoding, # note: `encode_video_frames` is a blocking call. Making it asynchronous shouldn't speedup encoding,
# since video encoding with ffmpeg is already using multithreading. # since video encoding with ffmpeg is already using multithreading.
encode_video_frames(tmp_imgs_dir, video_path, fps, overwrite=True) encode_video_frames(tmp_imgs_dir, video_path, fps, overwrite=True)
shutil.rmtree(tmp_imgs_dir)
logging.info("Concatenating episodes") logging.info("Concatenating episodes")
ep_dicts = [] ep_dicts = []