| import json |
| import os |
| import random |
|
|
| import numpy as np |
| import torch |
| import torchvision.transforms as transforms |
| from PIL import Image |
| from torch.utils.data.dataset import Dataset |
|
|
|
|
| class CC15M(Dataset): |
| def __init__( |
| self, |
| json_path, |
| video_folder=None, |
| resolution=512, |
| enable_bucket=False, |
| ): |
| print(f"loading annotations from {json_path} ...") |
| self.dataset = json.load(open(json_path, 'r')) |
| self.length = len(self.dataset) |
| print(f"data scale: {self.length}") |
| |
| self.enable_bucket = enable_bucket |
| self.video_folder = video_folder |
|
|
| resolution = tuple(resolution) if not isinstance(resolution, int) else (resolution, resolution) |
| self.pixel_transforms = transforms.Compose([ |
| transforms.Resize(resolution[0]), |
| transforms.CenterCrop(resolution), |
| transforms.ToTensor(), |
| transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), |
| ]) |
| |
| def get_batch(self, idx): |
| video_dict = self.dataset[idx] |
| video_id, name = video_dict['file_path'], video_dict['text'] |
|
|
| if self.video_folder is None: |
| video_dir = video_id |
| else: |
| video_dir = os.path.join(self.video_folder, video_id) |
|
|
| pixel_values = Image.open(video_dir).convert("RGB") |
| return pixel_values, name |
|
|
| def __len__(self): |
| return self.length |
|
|
| def __getitem__(self, idx): |
| while True: |
| try: |
| pixel_values, name = self.get_batch(idx) |
| break |
| except Exception as e: |
| print(e) |
| idx = random.randint(0, self.length-1) |
|
|
| if not self.enable_bucket: |
| pixel_values = self.pixel_transforms(pixel_values) |
| else: |
| pixel_values = np.array(pixel_values) |
|
|
| sample = dict(pixel_values=pixel_values, text=name) |
| return sample |
|
|
| if __name__ == "__main__": |
| dataset = CC15M( |
| csv_path="/mnt_wg/zhoumo.xjq/CCUtils/cc15m_add_index.json", |
| resolution=512, |
| ) |
| |
| dataloader = torch.utils.data.DataLoader(dataset, batch_size=4, num_workers=0,) |
| for idx, batch in enumerate(dataloader): |
| print(batch["pixel_values"].shape, len(batch["text"])) |
