Update emg2pose with s3 AWS CLI download

scripts/README.md

@@ -36,4 +36,4 @@ This guide provides commands to process raw EMG data into HDF5 format using slid
 | **NinaPro DB8** | Regression | 200 (0.1s) | 200 | `python scripts/db8.py --data_dir $DATA_PATH/ninapro/DB8/ --save_dir $DATA_PATH/ninapro/DB8/h5_100/ --seq_len 200 --stride 200` |
 | **NinaPro DB8** | Regression | 1000 (0.5s) | 1000 | `python scripts/db8.py --data_dir $DATA_PATH/ninapro/DB8/ --save_dir $DATA_PATH/ninapro/DB8/h5_500/ --seq_len 1000 --stride 1000` |
 
->Note: For DB5, we used the `--data-augment` flag to augment the training data by a factor of 3 (see `--augment-factor` in `scripts/db5.py`).
+>Note: For DB5, we used the `--data-augment` flag to augment the training data by a factor of 3 (see `--augment-factor` in `scripts/db5.py`).

scripts/emg2pose.py

@@ -1,5 +1,8 @@
 import os
 import gc
+import subprocess
+import tarfile
+import tempfile
 from pathlib import Path
 from typing import Tuple, List, Optional, Union, Dict, Any
 
@@ -11,6 +14,80 @@ from joblib import Parallel, delayed
 from scipy.signal import iirnotch
 from tqdm import tqdm
 
+
+EMG2POSE_S3_URI = "s3://fb-ctrl-oss/emg2pose/emg2pose_dataset.tar"
+
+
+def download_emg2pose(
+    data_dir: str,
+    archive_path: Optional[str] = None,
+    max_concurrent_requests: int = 64,
+    multipart_chunksize: str = "128MB",
+    io_chunksize: str = "64MB",
+    max_queue_size: int = 1000,
+    keep_archive: bool = False,
+) -> None:
+    """Download and extract the public EMG2Pose S3 archive.
+
+    AWS CLI S3 tuning is kept in a temporary config file so this script does
+    not modify the user's global AWS configuration. The public object is
+    downloaded without credentials and can be resumed by AWS CLI when a
+    partial file is present.
+    """
+    data_path = Path(data_dir).expanduser().resolve()
+    data_path.mkdir(parents=True, exist_ok=True)
+    archive = (
+        Path(archive_path).expanduser().resolve()
+        if archive_path
+        else data_path.parent / "emg2pose_dataset.tar"
+    )
+    archive.parent.mkdir(parents=True, exist_ok=True)
+
+    aws_config = (
+        "[default]\n"
+        "s3 =\n"
+        f"    max_concurrent_requests = {max_concurrent_requests}\n"
+        f"    max_queue_size = {max_queue_size}\n"
+        f"    multipart_chunksize = {multipart_chunksize}\n"
+        f"    io_chunksize = {io_chunksize}\n"
+    )
+
+    print(f"Downloading EMG2Pose to {archive}")
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".aws-config", delete=False) as config_file:
+        config_file.write(aws_config)
+        config_path = config_file.name
+
+    env = os.environ.copy()
+    env["AWS_CONFIG_FILE"] = config_path
+    try:
+        subprocess.run(
+            [
+                "aws",
+                "s3",
+                "cp",
+                EMG2POSE_S3_URI,
+                str(archive),
+                "--no-sign-request",
+            ],
+            check=True,
+            env=env,
+        )
+    finally:
+        os.unlink(config_path)
+
+    print(f"Extracting {archive} into {data_path}")
+    with tarfile.open(archive, "r") as tar:
+        root = data_path.resolve()
+        for member in tar.getmembers():
+            target = (data_path / member.name).resolve()
+            if target != root and root not in target.parents:
+                raise RuntimeError(f"Unsafe archive member: {member.name}")
+        tar.extractall(data_path)
+
+    if not keep_archive:
+        archive.unlink()
+        print("Removed downloaded archive")
+
 def sequence_to_seconds(seq_len: int, fs: float) -> float:
     """Converts a sequence length in samples to time in seconds.
 
@@ -135,6 +212,13 @@ def main():
     args = argparse.ArgumentParser(description="Process EMG data from DB5.")
     args.add_argument("--data_dir", type=str)
     args.add_argument("--save_dir", type=str)
+    args.add_argument("--download_data", action="store_true")
+    args.add_argument("--archive_path", type=str, default=None)
+    args.add_argument("--keep_archive", action="store_true")
+    args.add_argument("--aws_max_concurrent_requests", type=int, default=64)
+    args.add_argument("--aws_multipart_chunksize", type=str, default="128MB")
+    args.add_argument("--aws_io_chunksize", type=str, default="64MB")
+    args.add_argument("--aws_max_queue_size", type=int, default=1000)
     args.add_argument(
         "--seq_len", type=int, help="Size of the window in samples for segmentation."
     )
@@ -165,6 +249,17 @@ def main():
     save_dir = args.save_dir
     os.makedirs(save_dir, exist_ok=True)
 
+    if args.download_data:
+        download_emg2pose(
+            data_dir=data_dir,
+            archive_path=args.archive_path,
+            max_concurrent_requests=args.aws_max_concurrent_requests,
+            multipart_chunksize=args.aws_multipart_chunksize,
+            io_chunksize=args.aws_io_chunksize,
+            max_queue_size=args.aws_max_queue_size,
+            keep_archive=args.keep_archive,
+        )
+
     fs = 2000.0  # original sampling rate
     window_size, stride = args.seq_len, args.stride
 
@@ -230,4 +325,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()

scripts/requirements.txt

@@ -3,4 +3,5 @@ numpy
 scipy
 joblib
 tqdm
-pandas
+pandas
+awscli