InterDigitalInc / CompressAI

A PyTorch library and evaluation platform for end-to-end compression research
https://interdigitalinc.github.io/CompressAI/
BSD 3-Clause Clear License
1.21k stars 232 forks source link

Video datasets: Pregenerated, and Vimeo90K #272

Open YodaEmbedding opened 9 months ago

YodaEmbedding commented 9 months ago

Added mode="video" parameter to PreGeneratedMemmapDataset and Vimeo90kDataset.

To generate the pregenerated numpy datasets:

cd /data/datasets

mkdir -p vimeo90k
cd vimeo90k
wget http://data.csail.mit.edu/tofu/dataset/vimeo_triplet.zip
unzip vimeo_triplet.zip
wget http://data.csail.mit.edu/tofu/dataset/vimeo_septuplet.zip
unzip vimeo_septuplet.zip
cd ..

wget https://gist.githubusercontent.com/YodaEmbedding/8803d95de072f12b4ff14ffd2b5bd7e5/raw/9db2d7664984cc29b790f6b5268b971ec004c50c/generate_vimeo90k_npy_dataset.py

# NOTE: Ignore the first one if it already exists:
python generate_vimeo90k_npy_dataset.py --tuplet=3 --mode=image --indir="vimeo90k/vimeo_triplet" --outdir="vimeo90k/vimeo_triplet_npy"
python generate_vimeo90k_npy_dataset.py --tuplet=3 --mode=video --indir="vimeo90k/vimeo_triplet" --outdir="vimeo90k/vimeo_triplet_npy_video"
python generate_vimeo90k_npy_dataset.py --tuplet=7 --mode=image --indir="vimeo90k/vimeo_septuplet" --outdir="vimeo90k/vimeo_septuplet_npy"
python generate_vimeo90k_npy_dataset.py --tuplet=7 --mode=video --indir="vimeo90k/vimeo_septuplet" --outdir="vimeo90k/vimeo_septuplet_npy_video"

Example config (works only on latest CompressAI Trainer dev branch):

type: "PreGeneratedMemmapDataset"
config:
  root: "${paths.datasets}/vimeo90k/vimeo_septuplet_npy_video"
  split: "train"
  image_size: [ 256, 256, ]
  mode: "video"
  frames_per_sample: 7
loader:
  shuffle: True
  batch_size: 16
  num_workers: 2
settings:
  patch_size: [ 256, 256, ]
transforms:
  transform_frame:
    - "ToTensor": {}
  transform:
    - "RandomCrop": {size: "${.....settings.patch_size}"}
meta:
  name: "Vimeo-90K"
  identifier:
  num_samples: 64612
  steps_per_epoch:
fracape commented 7 months ago

Hi @YodaEmbedding, this has been marked as WIP for a while. Is it ready for merge?

YodaEmbedding commented 7 months ago

From my limited testing, it's probably good.

However, I did want to test it more thoroughly when working on Ujwal's project to confirm.

fracape commented 7 months ago

SG, let me know!