Create dataloader.py

dataloader.py

@@ -0,0 +1,56 @@
+import torch
+from datasets import load_dataset
+from torch.utils.data import DataLoader
+from transformers import BertTokenizer
+import decord
+import numpy as np
+from tqdm import tqdm
+
+
+FRAMES = 50
+H, W = 128, 128
+BATCH_SIZE = 8
+TEXT_MAX_LEN = 3000
+
+
+dataset = load_dataset("gaussalgo/webvid-10m", split="train")  # 10M samples
+
+
+tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+
+
+class VideoDataset(torch.utils.data.Dataset):
+    def __init__(self, dataset):
+        self.dataset = dataset
+        self.decord_ctx = decord.cpu(0)  # CPU decoding
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        item = self.dataset[idx]
+        
+
+        vr = decord.VideoReader(item["video_path"], ctx=self.decord_ctx)
+        frame_indices = np.linspace(0, len(vr)-1, FRAMES, dtype=int)
+        video = vr.get_batch(frame_indices).numpy()  # (FRAMES, H, W, 3)
+        video = torch.from_numpy(video).permute(3, 0, 1, 2).float()  # (3, FRAMES, H, W)
+        
+        
+        video = F.interpolate(video, size=(H, W), mode="bilinear")
+        video = (video / 255.0) * 2 - 1  # [-1, 1]
+        
+        
+        text = tokenizer(
+            item["caption"],
+            padding="max_length",
+            truncation=True,
+            max_length=TEXT_MAX_LEN,
+            return_tensors="pt"
+        ).input_ids.squeeze(0)
+        
+        return {"video": video, "text": text}
+
+# DataLoader
+dataset = VideoDataset(dataset)
+dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)