Update dataset

app.py

@@ -8,7 +8,7 @@ from tokenizers import Tokenizer
 from torch.utils.data import Dataset
 import albumentations as A
 from tqdm import tqdm
-
+from datasets import load_dataset
 from fourm.vq.vqvae import VQVAE
 from fourm.models.fm import FM
 from fourm.models.generate import (
@@ -28,7 +28,7 @@ IMG_SIZE = 224
 TOKENIZER_PATH = "./fourm/utils/tokenizer/trained/text_tokenizer_4m_wordpiece_30k.json"
 FM_MODEL_PATH = "EPFL-VILAB/4M-21_L"
 VQVAE_PATH = "EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224"
-IMAGE_DATASET_PATH = "/home/ubuntu/GIT_REPOS/ml-4m/data/custom_data/"
+IMAGE_DATASET_PATH = "./data"
 
 # Load models
 text_tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
@@ -61,25 +61,24 @@ schedule = build_chained_generation_schedules(
 sampler = GenerationSampler(fm_model)
 
 
-class ImageDataset(Dataset):
-    def __init__(self, path: str, img_sz=IMG_SIZE):
-        self.path = Path(path)
-        self.files = list(self.path.rglob("*"))
-        self.tfms = A.Compose(
-            [A.SmallestMaxSize(img_sz)])
+class HuggingFaceImageDataset(Dataset):
+    def __init__(self, dataset_name, split="train", img_sz=224):
+        self.dataset = load_dataset(dataset_name, split=split)
+        self.tfms = A.Compose([
+            A.SmallestMaxSize(img_sz)
+        ])
 
     def __len__(self):
-        return len(self.files)
+        return len(self.dataset)
 
     def __getitem__(self, idx):
-        img = Image.open(self.files[idx]).convert("RGB")
+        img = self.dataset[idx]['image']
         img = np.array(img)
         img = self.tfms(image=img)["image"]
         return Image.fromarray(img)
 
-
-dataset = ImageDataset(IMAGE_DATASET_PATH)
-
+# Usage
+dataset = HuggingFaceImageDataset("aroraaman/4m-21-demo")
 
 @torch.no_grad()
 def get_image_embeddings(dataset):