Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,935 Bytes
c382c9b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
import glob
import os
from PIL import Image
import random
import numpy as np
from torch import nn
from torchvision import transforms
from torch.utils import data as data
import torch.nn.functional as F
from .realesrgan import RealESRGAN_degradation
class PairedCaptionDataset(data.Dataset):
def __init__(
self,
root_folders=None,
tokenizer=None,
null_text_ratio=0.5,
# use_ram_encoder=False,
# use_gt_caption=False,
# caption_type = 'gt_caption',
):
super(PairedCaptionDataset, self).__init__()
self.null_text_ratio = null_text_ratio
self.lr_list = []
self.gt_list = []
self.tag_path_list = []
root_folders = root_folders.split(',')
for root_folder in root_folders:
lr_path = root_folder +'/sr_bicubic'
tag_path = root_folder +'/tag'
gt_path = root_folder +'/gt'
self.lr_list += glob.glob(os.path.join(lr_path, '*.png'))
self.gt_list += glob.glob(os.path.join(gt_path, '*.png'))
self.tag_path_list += glob.glob(os.path.join(tag_path, '*.txt'))
assert len(self.lr_list) == len(self.gt_list)
assert len(self.lr_list) == len(self.tag_path_list)
self.img_preproc = transforms.Compose([
transforms.ToTensor(),
])
ram_mean = [0.485, 0.456, 0.406]
ram_std = [0.229, 0.224, 0.225]
self.ram_normalize = transforms.Normalize(mean=ram_mean, std=ram_std)
self.tokenizer = tokenizer
def tokenize_caption(self, caption=""):
inputs = self.tokenizer(
caption, max_length=self.tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
)
return inputs.input_ids
def __getitem__(self, index):
gt_path = self.gt_list[index]
gt_img = Image.open(gt_path).convert('RGB')
gt_img = self.img_preproc(gt_img)
lq_path = self.lr_list[index]
lq_img = Image.open(lq_path).convert('RGB')
lq_img = self.img_preproc(lq_img)
if random.random() < self.null_text_ratio:
tag = ''
else:
tag_path = self.tag_path_list[index]
file = open(tag_path, 'r')
tag = file.read()
file.close()
example = dict()
example["conditioning_pixel_values"] = lq_img.squeeze(0)
example["pixel_values"] = gt_img.squeeze(0) * 2.0 - 1.0
example["input_ids"] = self.tokenize_caption(caption=tag).squeeze(0)
lq_img = lq_img.squeeze()
ram_values = F.interpolate(lq_img.unsqueeze(0), size=(384, 384), mode='bicubic')
ram_values = ram_values.clamp(0.0, 1.0)
example["ram_values"] = self.ram_normalize(ram_values.squeeze(0))
return example
def __len__(self):
return len(self.gt_list) |