|
import torch |
|
import os |
|
import torchvision.transforms as transforms |
|
import matplotlib.pyplot as plt |
|
import numpy as np |
|
|
|
from utils.utils import generate_mask |
|
|
|
|
|
class TrainDataset(torch.utils.data.Dataset): |
|
def __init__(self, data_path, transform=None, mults_amount=1): |
|
self.data = os.listdir(os.path.join(data_path, "color")) |
|
self.data_path = data_path |
|
self.transform = transform |
|
self.mults_amount = mults_amount |
|
|
|
self.ToTensor = transforms.ToTensor() |
|
|
|
def __len__(self): |
|
return len(self.data) |
|
|
|
def __getitem__(self, idx): |
|
image_name = self.data[idx] |
|
|
|
try: |
|
color_img = plt.imread(os.path.join(self.data_path, 'color', image_name)) |
|
except SyntaxError: |
|
print(f"Archivo {image_name} no es un PNG válido. Saltando...") |
|
return None |
|
|
|
if self.mults_amount > 1: |
|
mult_number = np.random.choice(range(self.mults_amount)) |
|
|
|
bw_name = ( |
|
image_name[: image_name.rfind(".")] + "_" + str(mult_number) + ".png" |
|
) |
|
dfm_name = ( |
|
image_name[: image_name.rfind(".")] |
|
+ "_" |
|
+ str(mult_number) |
|
+ "_dfm.png" |
|
) |
|
else: |
|
bw_name = self.data[idx] |
|
dfm_name = os.path.splitext(self.data[idx])[0] + "0_dfm.png" |
|
|
|
bw_img = np.expand_dims( |
|
plt.imread(os.path.join(self.data_path, "bw", bw_name)), 2 |
|
) |
|
dfm_img = np.expand_dims( |
|
plt.imread(os.path.join(self.data_path, "bw", dfm_name)), 2 |
|
) |
|
|
|
bw_img = np.concatenate([bw_img, dfm_img], axis=2) |
|
|
|
if self.transform: |
|
result = self.transform(image=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, mult_amount=1): |
|
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"))] |
|
self.data_path = data_path |
|
self.transform = transform |
|
self.mults_amount = mult_amount |
|
|
|
np.random.shuffle(self.color_data) |
|
|
|
self.ToTensor = transforms.ToTensor() |
|
|
|
def __len__(self): |
|
return len(self.data) |
|
|
|
def __getitem__(self, idx): |
|
color_img = plt.imread( |
|
os.path.join(self.data_path, "color", self.color_data[idx]) |
|
) |
|
|
|
image_name = self.data[idx] |
|
if self.mults_amount > 1: |
|
mult_number = np.random.choice(range(self.mults_amount)) |
|
|
|
bw_name = ( |
|
image_name[: image_name.rfind(".")] |
|
+ "_" |
|
+ str(self.mults_amount) |
|
+ ".png" |
|
) |
|
dfm_name = ( |
|
image_name[: image_name.rfind(".")] |
|
+ "_" |
|
+ str(self.mults_amount) |
|
+ "_dfm.png" |
|
) |
|
else: |
|
bw_name = self.data[idx] |
|
dfm_name = os.path.splitext(self.data[idx])[0] + "_dfm.png" |
|
|
|
bw_img = np.expand_dims( |
|
plt.imread(os.path.join(self.data_path, "real_manga", image_name)), 2 |
|
) |
|
dfm_img = np.expand_dims( |
|
plt.imread(os.path.join(self.data_path, "real_manga", dfm_name)), 2 |
|
) |
|
|
|
if self.transform: |
|
result = self.transform(image=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 |
|
|