YodaEmbedding · January 24, 2025 04:29 · Jan 24, 2025 · Jan 24, 2025 · Feb 26, 2024 · Feb 25, 2024
diff --git a/generate_vimeo90k_npy_dataset.py b/generate_vimeo90k_npy_dataset.py
@@ -1,13 +1,4 @@
-import argparse
-from pathlib import Path
-
-import numpy as np
-import torch
-from compressai.datasets import Vimeo90kDataset
-from torch.utils.data import DataLoader
-from torchvision import transforms
-
-MESSAGE = """
+"""
 Download and extract the Vimeo90k dataset first:
 
 mkdir -p vimeo90k
@@ -30,6 +21,16 @@
 If the mode is "video", all frames undergo the same transformation.
 """
 
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+from compressai.datasets import Vimeo90kDataset
+from torch.utils.data import DataLoader
+from torchvision import transforms
+
+
 PATCH_LENGTH = 256
 PATCH_SIZE = (PATCH_LENGTH, PATCH_LENGTH)
 
@@ -121,7 +122,7 @@ def parse_args():
 
 
 def main():
-    print(MESSAGE)
+    print(__doc__)
     args = parse_args()
     torch.manual_seed(args.seed)
 

diff --git a/generate_vimeo90k_npy_dataset.py b/generate_vimeo90k_npy_dataset.py
@@ -14,6 +14,8 @@
 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 ..
 
 Then, run one of the following:
@@ -96,7 +98,7 @@ def generate_npy_dataset(indir, outdir, split, tuplet, mode, epochs):
                 f"{split} | "
                 f"{epoch} / {epochs} epochs | "
                 f"{offset:6d} / {len(dataset)} items | "
-                f"{i:5d} / {len(loader)} batches | "
+                f"{i:6d} / {len(loader)} batches | "
                 # For ensuring that random output is stable:
                 f"checksum: {x.min():3.0f} {x.max():3.0f} {x.to(float).mean():3.0f}"
             )

diff --git a/generate_vimeo90k_npy_dataset.py b/generate_vimeo90k_npy_dataset.py
@@ -43,10 +43,18 @@ def get_dataset(dataset_path, split, tuplet, mode):
         if split == "train"
         else transforms.CenterCrop(PATCH_SIZE)
     )
+    chw_to_hwc = (
+        lambda x: x.permute(1, 2, 0)
+        if mode == "image"
+        else x.permute(0, -2, -1, -3)
+        if mode == "video"
+        else None
+    )
     transform = transforms.Compose(
         [
             crop,
-            lambda img: torch.from_numpy(np.array(img)),
+            # lambda img: torch.from_numpy(np.array(img)),
+            chw_to_hwc,
             # transforms.ToTensor(),  # NOTE: Converts HWC -> CHW.
         ]
     )
@@ -55,9 +63,11 @@ def get_dataset(dataset_path, split, tuplet, mode):
         transform=transform,
         split=split,
         tuplet=tuplet,
-        # mode is experimental. Old versions of CompressAI do not have
-        # this parameter, and behave as if mode="image".
+        # The following parameters are experimental.
+        # Old versions of CompressAI do not have these,
+        # and behave as if mode="image".
         mode=mode,
+        transform_frame=transforms.ToTensor(),  # NOTE: Converts HWC -> CHW.
     )
     loader = DataLoader(dataset, batch_size=16, shuffle=False, num_workers=8)
     return dataset, loader
@@ -70,12 +80,12 @@ def generate_npy_dataset(indir, outdir, split, tuplet, mode, epochs):
     out_filepath.parent.mkdir(exist_ok=True)
     print(f"Writing to {out_filepath}...")
 
-    x_out = np.memmap(
-        out_filepath,
-        dtype="uint8",
-        mode="w+",
-        shape=(epochs * len(dataset), *PATCH_SIZE, 3),
-    )
+    if mode == "image":
+        shape = (epochs * len(dataset), *PATCH_SIZE, 3)
+    elif mode == "video":
+        shape = (epochs * len(dataset), tuplet, *PATCH_SIZE, 3)
+
+    x_out = np.memmap(out_filepath, dtype="uint8", mode="w+", shape=shape)
 
     offset = 0
 

diff --git a/generate_vimeo90k_npy_dataset.py b/generate_vimeo90k_npy_dataset.py
@@ -74,7 +74,7 @@ def generate_npy_dataset(indir, outdir, split, tuplet, mode, epochs):
         out_filepath,
         dtype="uint8",
         mode="w+",
-        shape=(epochs * len(dataset), 3, *PATCH_SIZE),
+        shape=(epochs * len(dataset), *PATCH_SIZE, 3),
     )
 
     offset = 0
@@ -84,11 +84,11 @@ def generate_npy_dataset(indir, outdir, split, tuplet, mode, epochs):
             x = (x * 255).to(torch.uint8)
             print(
                 f"{split} | "
-                f"{epoch:2d} / {epochs} | "
-                f"{i:6d} / {len(loader)} | "
-                f"{offset:9d} | "
+                f"{epoch} / {epochs} epochs | "
+                f"{offset:6d} / {len(dataset)} items | "
+                f"{i:5d} / {len(loader)} batches | "
                 # For ensuring that random output is stable:
-                f"{x.min():3.0f} {x.max():3.0f} {x.to(float).mean():3.0f}"
+                f"checksum: {x.min():3.0f} {x.max():3.0f} {x.to(float).mean():3.0f}"
             )
             x_out[offset : offset + len(x)] = x.numpy()
             offset += len(x)

diff --git a/generate_vimeo90k_npy_dataset.py b/generate_vimeo90k_npy_dataset.py
@@ -63,7 +63,7 @@ def get_dataset(dataset_path, split, tuplet, mode):
     return dataset, loader
 
 
-def generate_npy_dataset(indir, outdir, split, tuplet, mode):
+def generate_npy_dataset(indir, outdir, split, tuplet, mode, epochs):
     dataset, loader = get_dataset(indir, split, tuplet, mode)
 
     out_filepath = Path(f"{outdir}/{FILENAMES[split]}.npy")
@@ -74,19 +74,24 @@ def generate_npy_dataset(indir, outdir, split, tuplet, mode):
         out_filepath,
         dtype="uint8",
         mode="w+",
-        shape=(len(dataset), 3, *PATCH_SIZE),
+        shape=(epochs * len(dataset), 3, *PATCH_SIZE),
     )
 
     offset = 0
 
-    for i, x in enumerate(loader):
-        x = (x * 255).to(torch.uint8)
-        print(
-            f"{i:4d} / {len(loader)} {offset:5.0f} | "
-            f"{x.min():3.0f} {x.max():3.0f} {x.to(float).mean():3.0f}"
-        )
-        x_out[offset : offset + len(x)] = x.numpy()
-        offset += len(x)
+    for epoch in range(epochs):
+        for i, x in enumerate(loader):
+            x = (x * 255).to(torch.uint8)
+            print(
+                f"{split} | "
+                f"{epoch:2d} / {epochs} | "
+                f"{i:6d} / {len(loader)} | "
+                f"{offset:9d} | "
+                # For ensuring that random output is stable:
+                f"{x.min():3.0f} {x.max():3.0f} {x.to(float).mean():3.0f}"
+            )
+            x_out[offset : offset + len(x)] = x.numpy()
+            offset += len(x)
 
     x_out.flush()
     del x_out
@@ -99,6 +104,7 @@ def parse_args():
     parser.add_argument("--tuplet", type=int, default=3)
     parser.add_argument("--mode", default="image", choices=["image", "video"])
     parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument("--epochs", type=int, default=1)
     return parser.parse_args()
 
 
@@ -114,6 +120,7 @@ def main():
             split=split,
             tuplet=args.tuplet,
             mode=args.mode,
+            epochs=args.epochs,
         )
 
 

diff --git a/generate_vimeo90k_npy_dataset.py b/generate_vimeo90k_npy_dataset.py
@@ -55,6 +55,9 @@ def get_dataset(dataset_path, split, tuplet, mode):
         transform=transform,
         split=split,
         tuplet=tuplet,
+        # mode is experimental. Old versions of CompressAI do not have
+        # this parameter, and behave as if mode="image".
+        mode=mode,
     )
     loader = DataLoader(dataset, batch_size=16, shuffle=False, num_workers=8)
     return dataset, loader

diff --git a/generate_vimeo90k_npy_dataset.py b/generate_vimeo90k_npy_dataset.py
@@ -16,10 +16,16 @@
 unzip vimeo_triplet.zip
 cd ..
 
-Then, run:
+Then, run one of the following:
 
-python generate_vimeo90k_npy_dataset.py --dataset="vimeo90k/vimeo_triplet" --tuplet=3
+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"
 
+If the mode is "image", each frame is treated separately, and may
+undergo different transformations.
+If the mode is "video", all frames undergo the same transformation.
 """
 
 PATCH_LENGTH = 256
@@ -31,33 +37,33 @@
 }
 
 
-def get_dataset(dataset_path, split, tuplet):
+def get_dataset(dataset_path, split, tuplet, mode):
     crop = (
         transforms.RandomCrop(PATCH_SIZE)
         if split == "train"
         else transforms.CenterCrop(PATCH_SIZE)
     )
+    transform = transforms.Compose(
+        [
+            crop,
+            lambda img: torch.from_numpy(np.array(img)),
+            # transforms.ToTensor(),  # NOTE: Converts HWC -> CHW.
+        ]
+    )
     dataset = Vimeo90kDataset(
         root=dataset_path,
-        transform=transforms.Compose(
-            [
-                crop,
-                transforms.ToTensor(),
-            ]
-        ),
+        transform=transform,
         split=split,
         tuplet=tuplet,
     )
     loader = DataLoader(dataset, batch_size=16, shuffle=False, num_workers=8)
     return dataset, loader
 
 
-def generate_npy_dataset(dataset_path, split, tuplet):
-    dataset, loader = get_dataset(dataset_path, split, tuplet)
+def generate_npy_dataset(indir, outdir, split, tuplet, mode):
+    dataset, loader = get_dataset(indir, split, tuplet, mode)
 
-    out_filepath = Path(
-        f"{str(Path(dataset_path)).rstrip('/')}_npy/{FILENAMES[split]}.npy"
-    )
+    out_filepath = Path(f"{outdir}/{FILENAMES[split]}.npy")
     out_filepath.parent.mkdir(exist_ok=True)
     print(f"Writing to {out_filepath}...")
 
@@ -85,17 +91,27 @@ def generate_npy_dataset(dataset_path, split, tuplet):
 
 def parse_args():
     parser = argparse.ArgumentParser(description="Generate Vimeo90k dataset")
-    parser.add_argument("--dataset", default="vimeo90k/vimeo_triplet")
+    parser.add_argument("--indir", default="vimeo90k/vimeo_triplet")
+    parser.add_argument("--outdir", default="vimeo90k/vimeo_triplet_npy")
     parser.add_argument("--tuplet", type=int, default=3)
+    parser.add_argument("--mode", default="image", choices=["image", "video"])
+    parser.add_argument("--seed", type=int, default=1234)
     return parser.parse_args()
 
 
 def main():
     print(MESSAGE)
     args = parse_args()
-    torch.manual_seed(1234)
-    generate_npy_dataset(dataset_path=args.dataset, split="train", tuplet=args.tuplet)
-    generate_npy_dataset(dataset_path=args.dataset, split="valid", tuplet=args.tuplet)
+    torch.manual_seed(args.seed)
+
+    for split in ["train", "valid"]:
+        generate_npy_dataset(
+            indir=args.indir,
+            outdir=args.outdir,
+            split=split,
+            tuplet=args.tuplet,
+            mode=args.mode,
+        )
 
 
 if __name__ == "__main__":

diff --git a/generate_vimeo90k_npy_dataset.py b/generate_vimeo90k_npy_dataset.py
@@ -25,6 +25,11 @@
 PATCH_LENGTH = 256
 PATCH_SIZE = (PATCH_LENGTH, PATCH_LENGTH)
 
+FILENAMES = {
+    "train": "training",
+    "valid": "validation",
+}
+
 
 def get_dataset(dataset_path, split, tuplet):
     crop = (
@@ -50,7 +55,9 @@ def get_dataset(dataset_path, split, tuplet):
 def generate_npy_dataset(dataset_path, split, tuplet):
     dataset, loader = get_dataset(dataset_path, split, tuplet)
 
-    out_filepath = Path(f"{str(Path(dataset_path)).rstrip('/')}_npy/{split}.npy")
+    out_filepath = Path(
+        f"{str(Path(dataset_path)).rstrip('/')}_npy/{FILENAMES[split]}.npy"
+    )
     out_filepath.parent.mkdir(exist_ok=True)
     print(f"Writing to {out_filepath}...")
 

diff --git a/generate_vimeo90k_npy_dataset.py b/generate_vimeo90k_npy_dataset.py
@@ -1,13 +1,12 @@
+import argparse
 from pathlib import Path
 
 import numpy as np
 import torch
-
+from compressai.datasets import Vimeo90kDataset
 from torch.utils.data import DataLoader
 from torchvision import transforms
 
-from compressai.datasets import Vimeo90kDataset
-
 MESSAGE = """
 Download and extract the Vimeo90k dataset first:
 
@@ -16,42 +15,50 @@
 wget http://data.csail.mit.edu/tofu/dataset/vimeo_triplet.zip
 unzip vimeo_triplet.zip
 cd ..
+
+Then, run:
+
+python generate_vimeo90k_npy_dataset.py --dataset="vimeo90k/vimeo_triplet" --tuplet=3
+
 """
 
-DATASET_PATH = "vimeo90k/vimeo_triplet"
 PATCH_LENGTH = 256
 PATCH_SIZE = (PATCH_LENGTH, PATCH_LENGTH)
 
 
-def get_dataset(split):
+def get_dataset(dataset_path, split, tuplet):
     crop = (
         transforms.RandomCrop(PATCH_SIZE)
         if split == "train"
         else transforms.CenterCrop(PATCH_SIZE)
     )
     dataset = Vimeo90kDataset(
-        root=DATASET_PATH,
+        root=dataset_path,
         transform=transforms.Compose(
             [
                 crop,
                 transforms.ToTensor(),
             ]
         ),
         split=split,
+        tuplet=tuplet,
     )
     loader = DataLoader(dataset, batch_size=16, shuffle=False, num_workers=8)
     return dataset, loader
 
 
-def generate_npy_dataset(split):
-    dataset, loader = get_dataset(split)
+def generate_npy_dataset(dataset_path, split, tuplet):
+    dataset, loader = get_dataset(dataset_path, split, tuplet)
 
-    out_filepath = Path(DATASET_PATH).parent / f"vimeo_triplet_npy/{split}.npy"
+    out_filepath = Path(f"{str(Path(dataset_path)).rstrip('/')}_npy/{split}.npy")
     out_filepath.parent.mkdir(exist_ok=True)
     print(f"Writing to {out_filepath}...")
 
     x_out = np.memmap(
-        out_filepath, dtype="uint8", mode="w+", shape=(len(dataset), 3, *PATCH_SIZE)
+        out_filepath,
+        dtype="uint8",
+        mode="w+",
+        shape=(len(dataset), 3, *PATCH_SIZE),
     )
 
     offset = 0
@@ -69,11 +76,19 @@ def generate_npy_dataset(split):
     del x_out
 
 
+def parse_args():
+    parser = argparse.ArgumentParser(description="Generate Vimeo90k dataset")
+    parser.add_argument("--dataset", default="vimeo90k/vimeo_triplet")
+    parser.add_argument("--tuplet", type=int, default=3)
+    return parser.parse_args()
+
+
 def main():
     print(MESSAGE)
+    args = parse_args()
     torch.manual_seed(1234)
-    generate_npy_dataset(split="train")
-    generate_npy_dataset(split="valid")
+    generate_npy_dataset(dataset_path=args.dataset, split="train", tuplet=args.tuplet)
+    generate_npy_dataset(dataset_path=args.dataset, split="valid", tuplet=args.tuplet)
 
 
 if __name__ == "__main__":

diff --git a/generate_vimeo90k_npy_dataset.py b/generate_vimeo90k_npy_dataset.py
@@ -57,7 +57,11 @@ def generate_npy_dataset(split):
     offset = 0
 
     for i, x in enumerate(loader):
-        print(f"{i} / {len(loader)}    {offset}    {x.shape}")
+        x = (x * 255).to(torch.uint8)
+        print(
+            f"{i:4d} / {len(loader)} {offset:5.0f} | "
+            f"{x.min():3.0f} {x.max():3.0f} {x.to(float).mean():3.0f}"
+        )
         x_out[offset : offset + len(x)] = x.numpy()
         offset += len(x)
 

diff --git a/generate_vimeo90k_npy_dataset.py b/generate_vimeo90k_npy_dataset.py
@@ -1,6 +1,7 @@
 from pathlib import Path
 
 import numpy as np
+import torch
 
 from torch.utils.data import DataLoader
 from torchvision import transforms
@@ -66,6 +67,7 @@ def generate_npy_dataset(split):
 
 def main():
     print(MESSAGE)
+    torch.manual_seed(1234)
     generate_npy_dataset(split="train")
     generate_npy_dataset(split="valid")
 

diff --git a/generate_vimeo90k_npy_dataset.py b/generate_vimeo90k_npy_dataset.py
@@ -0,0 +1,74 @@
+from pathlib import Path
+
+import numpy as np
+
+from torch.utils.data import DataLoader
+from torchvision import transforms
+
+from compressai.datasets import Vimeo90kDataset
+
+MESSAGE = """
+Download and extract the Vimeo90k dataset first:
+
+mkdir -p vimeo90k
+cd vimeo90k
+wget http://data.csail.mit.edu/tofu/dataset/vimeo_triplet.zip
+unzip vimeo_triplet.zip
+cd ..
+"""
+
+DATASET_PATH = "vimeo90k/vimeo_triplet"
+PATCH_LENGTH = 256
+PATCH_SIZE = (PATCH_LENGTH, PATCH_LENGTH)
+
+
+def get_dataset(split):
+    crop = (
+        transforms.RandomCrop(PATCH_SIZE)
+        if split == "train"
+        else transforms.CenterCrop(PATCH_SIZE)
+    )
+    dataset = Vimeo90kDataset(
+        root=DATASET_PATH,
+        transform=transforms.Compose(
+            [
+                crop,
+                transforms.ToTensor(),
+            ]
+        ),
+        split=split,
+    )
+    loader = DataLoader(dataset, batch_size=16, shuffle=False, num_workers=8)
+    return dataset, loader
+
+
+def generate_npy_dataset(split):
+    dataset, loader = get_dataset(split)
+
+    out_filepath = Path(DATASET_PATH).parent / f"vimeo_triplet_npy/{split}.npy"
+    out_filepath.parent.mkdir(exist_ok=True)
+    print(f"Writing to {out_filepath}...")
+
+    x_out = np.memmap(
+        out_filepath, dtype="uint8", mode="w+", shape=(len(dataset), 3, *PATCH_SIZE)
+    )
+
+    offset = 0
+
+    for i, x in enumerate(loader):
+        print(f"{i} / {len(loader)}    {offset}    {x.shape}")
+        x_out[offset : offset + len(x)] = x.numpy()
+        offset += len(x)
+
+    x_out.flush()
+    del x_out
+
+
+def main():
+    print(MESSAGE)
+    generate_npy_dataset(split="train")
+    generate_npy_dataset(split="valid")
+
+
+if __name__ == "__main__":
+    main()