|
|
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 |
|
|
|
|
|
class ImageEditDataset(Dataset): |
|
|
def __init__( |
|
|
self, |
|
|
ann_path, data_root=None, |
|
|
image_sample_size=512, |
|
|
text_drop_ratio=0.1, |
|
|
enable_bucket=False, |
|
|
enable_inpaint=False, |
|
|
return_file_name=False, |
|
|
): |
|
|
|
|
|
print(f"loading annotations from {ann_path} ...") |
|
|
if ann_path.endswith('.csv'): |
|
|
with open(ann_path, 'r') as csvfile: |
|
|
dataset = list(csv.DictReader(csvfile)) |
|
|
elif ann_path.endswith('.json'): |
|
|
dataset = json.load(open(ann_path)) |
|
|
|
|
|
self.data_root = data_root |
|
|
self.dataset = dataset |
|
|
|
|
|
self.length = len(self.dataset) |
|
|
print(f"data scale: {self.length}") |
|
|
|
|
|
self.enable_bucket = enable_bucket |
|
|
self.text_drop_ratio = text_drop_ratio |
|
|
self.enable_inpaint = enable_inpaint |
|
|
self.return_file_name = return_file_name |
|
|
|
|
|
|
|
|
self.image_sample_size = tuple(image_sample_size) if not isinstance(image_sample_size, int) else (image_sample_size, image_sample_size) |
|
|
self.image_transforms = transforms.Compose([ |
|
|
transforms.Resize(min(self.image_sample_size)), |
|
|
transforms.CenterCrop(self.image_sample_size), |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize([0.5, 0.5, 0.5],[0.5, 0.5, 0.5]) |
|
|
]) |
|
|
|
|
|
def get_batch(self, idx): |
|
|
data_info = self.dataset[idx % len(self.dataset)] |
|
|
|
|
|
image_path, text = data_info['file_path'], data_info['text'] |
|
|
if self.data_root is not None: |
|
|
image_path = os.path.join(self.data_root, image_path) |
|
|
image = Image.open(image_path).convert('RGB') |
|
|
|
|
|
if not self.enable_bucket: |
|
|
raise ValueError("Not enable_bucket is not supported now. ") |
|
|
else: |
|
|
image = np.expand_dims(np.array(image), 0) |
|
|
|
|
|
source_image_path = data_info.get('source_file_path', []) |
|
|
source_image = [] |
|
|
if isinstance(source_image_path, list): |
|
|
for _source_image_path in source_image_path: |
|
|
if self.data_root is not None: |
|
|
_source_image_path = os.path.join(self.data_root, _source_image_path) |
|
|
_source_image = Image.open(_source_image_path).convert('RGB') |
|
|
source_image.append(_source_image) |
|
|
else: |
|
|
if self.data_root is not None: |
|
|
_source_image_path = os.path.join(self.data_root, source_image_path) |
|
|
_source_image = Image.open(_source_image_path).convert('RGB') |
|
|
source_image.append(_source_image) |
|
|
|
|
|
if not self.enable_bucket: |
|
|
raise ValueError("Not enable_bucket is not supported now. ") |
|
|
else: |
|
|
source_image = [np.array(_source_image) for _source_image in source_image] |
|
|
|
|
|
if random.random() < self.text_drop_ratio: |
|
|
text = '' |
|
|
return image, source_image, text, 'image', image_path |
|
|
|
|
|
def __len__(self): |
|
|
return self.length |
|
|
|
|
|
def __getitem__(self, idx): |
|
|
data_info = self.dataset[idx % len(self.dataset)] |
|
|
data_type = data_info.get('type', 'image') |
|
|
while True: |
|
|
sample = {} |
|
|
try: |
|
|
data_info_local = self.dataset[idx % len(self.dataset)] |
|
|
data_type_local = data_info_local.get('type', 'image') |
|
|
if data_type_local != data_type: |
|
|
raise ValueError("data_type_local != data_type") |
|
|
|
|
|
pixel_values, source_pixel_values, name, data_type, file_path = self.get_batch(idx) |
|
|
sample["pixel_values"] = pixel_values |
|
|
sample["source_pixel_values"] = source_pixel_values |
|
|
sample["text"] = name |
|
|
sample["data_type"] = data_type |
|
|
sample["idx"] = idx |
|
|
if self.return_file_name: |
|
|
sample["file_name"] = os.path.basename(file_path) |
|
|
|
|
|
if len(sample) > 0: |
|
|
break |
|
|
except Exception as e: |
|
|
print(e, self.dataset[idx % len(self.dataset)]) |
|
|
idx = random.randint(0, self.length-1) |
|
|
|
|
|
if self.enable_inpaint and not self.enable_bucket: |
|
|
mask = get_random_mask(pixel_values.size()) |
|
|
mask_pixel_values = pixel_values * (1 - mask) + torch.ones_like(pixel_values) * -1 * mask |
|
|
sample["mask_pixel_values"] = mask_pixel_values |
|
|
sample["mask"] = mask |
|
|
|
|
|
clip_pixel_values = sample["pixel_values"][0].permute(1, 2, 0).contiguous() |
|
|
clip_pixel_values = (clip_pixel_values * 0.5 + 0.5) * 255 |
|
|
sample["clip_pixel_values"] = clip_pixel_values |
|
|
|
|
|
return sample |
|
|
|
|
|
if __name__ == "__main__": |
|
|
dataset = CC15M( |
|
|
csv_path="./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"])) |
