John6666's picture
Upload 49 files
f1d6334 verified
raw
history blame
1.55 kB
import os
import pandas as pd
import numpy as np
from PIL import Image
import torch
from torch.utils.data import Dataset, DataLoader
import json
import random
def c_crop(image):
width, height = image.size
new_size = min(width, height)
left = (width - new_size) / 2
top = (height - new_size) / 2
right = (width + new_size) / 2
bottom = (height + new_size) / 2
return image.crop((left, top, right, bottom))
class CustomImageDataset(Dataset):
def __init__(self, img_dir, img_size=512):
self.images = [os.path.join(img_dir, i) for i in os.listdir(img_dir) if '.jpg' in i or '.png' in i]
self.images.sort()
self.img_size = img_size
def __len__(self):
return len(self.images)
def __getitem__(self, idx):
try:
img = Image.open(self.images[idx])
img = c_crop(img)
img = img.resize((self.img_size, self.img_size))
img = torch.from_numpy((np.array(img) / 127.5) - 1)
img = img.permute(2, 0, 1)
json_path = self.images[idx].split('.')[0] + '.json'
prompt = json.load(open(json_path))['caption']
return img, prompt
except Exception as e:
print(e)
return self.__getitem__(random.randint(0, len(self.images) - 1))
def loader(train_batch_size, num_workers, **args):
dataset = CustomImageDataset(**args)
return DataLoader(dataset, batch_size=train_batch_size, num_workers=num_workers, shuffle=True)