Downloading and Preprocessing Datasets

Import necesary modules

from tqdm.notebook import tqdm

import os

from torch_skeleton.datasets import UCLA, Apply, DiskCache
import torch_skeleton.transforms as T

/home/docs/checkouts/readthedocs.org/user_builds/torch-skeleton/envs/v0.1.2/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

Download and Preprocess Dataset

preprocess = T.Compose(
    [
        T.CenterJoint(joint_id=1, all=False),
        T.ParallelBone(first_id=0, second_id=1, axis=2),
        T.ParallelBone(first_id=8, second_id=4, axis=0),
    ]
)

dataset = UCLA(root="data", split="train", transform=preprocess)

print(os.listdir("data/NW-UCLA/all_sqe")[:5])

# apply preprocessing at creation
dataset = UCLA(root="data", transform=preprocess)

x, y = dataset[0]

m, t, v, c = x.shape
print(f"returns x with M={m}, T={t}, V={v}, C={c} {x.shape}")

# or apply them after creation
dataset = UCLA(root="data")
dataset = Apply(dataset, transform=preprocess)

Downloading all_sqe.zip

['a02_s05_e00_v01.json', 'a03_s01_e04_v01.json', 'a11_s01_e01_v03.json', 'a12_s08_e03_v03.json', 'a05_s10_e00_v01.json']
all_sqe.zip exists, skipping download
returns x with M=1, T=44, V=20, C=3 (1, 44, 20, 3)
all_sqe.zip exists, skipping download

Cache Preprocessed Dataset to Disk

cache = DiskCache(dataset, root="/tmp/ucla")


def list_temp_dir():
    tmp_dir = os.listdir("/tmp/ucla")[0]
    os.listdir(os.path.join("/tmp/ucla", tmp_dir))[:10]


list_temp_dir()
cache[0]
list_temp_dir()

for x, y in cache:
    pass

list_temp_dir()

Add augmentations for training

dataset = Apply(
    cache,
    transform=T.Compose(
        [
            T.SampleFrames(num_frames=20),
            T.RandomRotate(degrees=17),
            T.PadFrames(max_frames=20),
        ]
    ),
)