p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] pyproject.toml (plaintext)
[project]
name = "jasmine"
version = "0.1.0"
requires-python = ">=3.11"
dependencies = [
  "dm-pix>=0.4.3",
  "einops>=0.8.0",
  "flax>=0.10.7",
  "jax[cuda12]>=0.6.2",
  "optax>=0.2.3",
  "tyro>=0.8.5",
  "wandb>=0.17.4",
  "grain>=0.2.10",
  "array-record>=0.7.2",
  "pre-commit>=4.2.0",
  "pillow>=11.3.0",
]
[content] data/pyproject.toml (plaintext)
[project]
name = "jasmine-data"
version = "0.1.0"
requires-python = "==3.10"
dependencies = [
  "array-record>=0.7.2",
  "numpy",
  "procgen>=0.10.7",
  "hf-transfer==0.1.9",
  "huggingface-hub[cli]>=0.34.3",
  "ffmpeg-python==0.2.0",
  "pillow>=11.3.0",
  "tqdm>=4.67.1",
]
[tab] jasmine/utils/preprocess.py (python)
import einops
import jax
import jax.numpy as jnp


def patchify(videos: jax.Array, size: int) -> jax.Array:
    B, T, H, W, C = videos.shape
    x = jnp.pad(videos, ((0, 0), (0, 0), (0, -H % size), (0, -W % size), (0, 0)))
    return einops.rearrange(
        x, "b t (hn hp) (wn wp) c -> b t (hn wn) (hp wp c)", hp=size, wp=size
    )


def unpatchify(patches: jax.Array, size: int, h_out: int, w_out: int) -> jax.Array:
    h_pad = -h_out % size
    hn = (h_out + h_pad) // size
    x = einops.rearrange(
        patches,
        "b t (hn wn) (hp wp c) -> b t (hn hp) (wn wp) c",
        hp=size,
        wp=size,
        hn=hn,
    )
    return x[:, :, :h_out, :w_out]