ControlNet/tutorial_dataset.py at main · chokami/ControlNet · GitHub

39 lines (27 loc) · 1.08 KB

import json

import cv2

import numpy as np

from torch.utils.data import Dataset

class MyDataset(Dataset):

def __init__(self):

self.data = []

with open('./training/fill50k/prompt.json', 'rt') as f:

for line in f:

self.data.append(json.loads(line))

def __len__(self):

return len(self.data)

def __getitem__(self, idx):

item = self.data[idx]

source_filename = item['source']

target_filename = item['target']

prompt = item['prompt']

source = cv2.imread('./training/fill50k/' + source_filename)

target = cv2.imread('./training/fill50k/' + target_filename)

# Do not forget that OpenCV read images in BGR order.

source = cv2.cvtColor(source, cv2.COLOR_BGR2RGB)

target = cv2.cvtColor(target, cv2.COLOR_BGR2RGB)

# Normalize source images to [0, 1].

source = source.astype(np.float32) / 255.0

# Normalize target images to [-1, 1].

target = (target.astype(np.float32) / 127.5) - 1.0

return dict(jpg=target, txt=prompt, hint=source)

Expand file tree