Milcho/ControlNet-Guidance

1import json2import cv23import numpy as np4 5from torch.utils.data import Dataset6 7 8class MyDataset(Dataset):9    def __init__(self):10        self.data = []11        with open('./training/fill50k/prompt.json', 'rt') as f:12            for line in f:13                self.data.append(json.loads(line))14 15    def __len__(self):16        return len(self.data)17 18    def __getitem__(self, idx):19        item = self.data[idx]20 21        source_filename = item['source']22        target_filename = item['target']23        prompt = item['prompt']24 25        source = cv2.imread('./training/fill50k/' + source_filename)26        target = cv2.imread('./training/fill50k/' + target_filename)27 28        # Do not forget that OpenCV read images in BGR order.29        source = cv2.cvtColor(source, cv2.COLOR_BGR2RGB)30        target = cv2.cvtColor(target, cv2.COLOR_BGR2RGB)31 32        # Normalize source images to [0, 1].33        source = source.astype(np.float32) / 255.034 35        # Normalize target images to [-1, 1].36        target = (target.astype(np.float32) / 127.5) - 1.037 38        return dict(jpg=target, txt=prompt, hint=source)39 40