File size: 3,536 Bytes
52acd7f c17106f c51eaca 45c435a 09800e1 45c435a c17106f 9a6de71 45c435a c17106f 45c435a c17106f 45c435a 3422c59 09800e1 c17106f 9a6de71 c17106f 45c435a c17106f 45c435a c17106f 0c68476 c17106f 0c68476 c17106f 0c68476 c17106f c51eaca 0c68476 c17106f 0c68476 c17106f 0c68476 c17106f 7c6d953 c340733 c17106f 9a6de71 c340733 c17106f c340733 c17106f c340733 c17106f 9a6de71 c17106f c340733 c17106f c340733 c17106f c340733 c17106f c340733 c17106f c340733 c17106f 7c6d953 |
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 96 97 98 99 |
import torch
import os
from PIL import Image
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
from utils.utils import generate_mask
from PIL import Image
class TrainDataset(torch.utils.data.Dataset):
def __init__(self, data_path, transform=None):
self.data = os.listdir(os.path.join(data_path, 'color'))
self.data_path = data_path
self.transform = transform
self.ToTensor = transforms.ToTensor()
# Directorio para guardar las imágenes en blanco y negro
self.bw_directory = os.path.join(data_path, 'bw')
if not os.path.exists(self.bw_directory):
os.makedirs(self.bw_directory)
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
image_name = self.data[idx]
color_img = Image.open(os.path.join(self.data_path, 'color', image_name)).convert('RGB')
bw_name = "bw_" + image_name
dfm_name = 'dfm_' + image_name
bw_img = np.expand_dims(np.array(Image.open(os.path.join(self.bw_directory, bw_name)).convert('L')), 2)
dfm_img = np.expand_dims(np.array(Image.open(os.path.join(self.bw_directory, dfm_name)).convert('L')), 2)
bw_img = np.concatenate([bw_img, dfm_img], axis=2)
if self.transform:
result = self.transform(image=np.array(color_img), mask=bw_img)
color_img = result['image']
bw_img = result['mask']
dfm_img = bw_img[:, :, 1]
bw_img = bw_img[:, :, 0]
color_img = self.ToTensor(color_img)
bw_img = self.ToTensor(bw_img)
dfm_img = self.ToTensor(dfm_img)
color_img = (color_img - 0.5) / 0.5
mask = generate_mask(bw_img.shape[1], bw_img.shape[2])
hint = torch.cat((color_img * mask, mask), 0)
return bw_img, color_img, hint, dfm_img
class FineTuningDataset(torch.utils.data.Dataset):
def __init__(self, data_path, transform=None):
self.data = [x for x in os.listdir(os.path.join(data_path, 'real_manga')) if x.find('dfm_') == -1]
self.color_data = [x for x in os.listdir(os.path.join(data_path, 'color'))] * 100
self.data_path = data_path
self.transform = transform
# Directorio para guardar las imágenes en blanco y negro
self.bw_directory = os.path.join(data_path, 'bw')
if not os.path.exists(self.bw_directory):
os.makedirs(self.bw_directory)
np.random.shuffle(self.color_data)
self.ToTensor = transforms.ToTensor()
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
color_img = Image.open(os.path.join(self.data_path, 'color', self.color_data[idx])).convert('RGB')
bw_name = "bw_" + self.data[idx]
dfm_name = "dfm_" + self.data[idx]
bw_img = np.expand_dims(np.array(Image.open(os.path.join(self.bw_directory, bw_name)).convert('L')), 2)
dfm_img = np.expand_dims(np.array(Image.open(os.path.join(self.bw_directory, dfm_name)).convert('L')), 2)
if self.transform:
result = self.transform(image=np.array(color_img))
color_img = result['image']
result = self.transform(image=bw_img, mask=dfm_img)
bw_img = result['image']
dfm_img = result['mask']
color_img = self.ToTensor(color_img)
bw_img = self.ToTensor(bw_img)
dfm_img = self.ToTensor(dfm_img)
color_img = (color_img - 0.5) / 0.5
return bw_img, dfm_img, color_img
|