2025-11-12 00:59:35 +08:00

61 lines
2.1 KiB
Bash

BEGIN_TIME=$(date +%s)
CONFIG_FILE="input/config.json"
echo "CONFIG_FILE_PATH: $CONFIG_FILE"
# Read values directly from the config.json using python - no more nested key error by using a helper script
TASK_ID=$(python3 read_json.py "$CONFIG_FILE" "task_id")
DATA_DIR=$(python3 read_json.py "$CONFIG_FILE" "data_dir")
OUTPUT_DIR=$(python3 read_json.py "$CONFIG_FILE" "output_dir")
EPISODE_NUM=$(python3 read_json.py "$CONFIG_FILE" "episode_num")
GPU=$(python3 read_json.py "$CONFIG_FILE" "gpu")
T5_PATH="/weights/t5-v1_1-xxl"
NO_LANGUAGE=$(python3 read_json.py "$CONFIG_FILE" "no_language")
# For the camera keys, extract them in a way that avoids the error about 'images_info.key.*' not found
CAM_HIGH_KEY=$(python3 -c "import json; print(json.load(open('$CONFIG_FILE'))['images_info']['key'].get('cam_high', ''))")
CAM_RIGHT_WRIST_KEY=$(python3 -c "import json; print(json.load(open('$CONFIG_FILE'))['images_info']['key'].get('cam_right_wrist', ''))")
# create output path
if [ ! -d "$OUTPUT_DIR/$TASK_ID" ]; then
mkdir -p "$OUTPUT_DIR/$TASK_ID"
echo "Created output directory: $OUTPUT_DIR/$TASK_ID"
else
echo "Output directory already exists: $OUTPUT_DIR/$TASK_ID"
fi
if [ "$NO_LANGUAGE" = "true" ]; then
python3 lerobot2rdt.py \
--data_dir $DATA_DIR \
--output_dir $OUTPUT_DIR/$TASK_ID \
--episode_num $EPISODE_NUM \
--gpu $GPU \
--t5_path $T5_PATH \
--cam_high_key $CAM_HIGH_KEY \
--cam_right_wrist_key $CAM_RIGHT_WRIST_KEY \
--no_language
status=$?
else
python3 lerobot2rdt.py \
--data_dir $DATA_DIR \
--output_dir $OUTPUT_DIR/$TASK_ID \
--episode_num $EPISODE_NUM \
--gpu $GPU \
--t5_path $T5_PATH \
--cam_high_key $CAM_HIGH_KEY \
--cam_right_wrist_key $CAM_RIGHT_WRIST_KEY
status=$?
fi
END_TIME=$(date +%s)
echo "END_TIME: $END_TIME"
echo "TOTAL_TIME: $((END_TIME - BEGIN_TIME))"
if [ $status -eq 0 ]; then
python3 generate_output.py $CONFIG_FILE $((END_TIME - BEGIN_TIME))
else
echo "lerobot2rdt.py exited with status $status, skipping generate_output.py"
fi