Hi,
Could you please help me find the mismatch between the CustomDataset and the model?
Thanks!
Code below:
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
class CustomDataset(Dataset):
def __init__(self, path,start_idx, end_idx):
self.data = np.load(path)
self.data = self.data[start_idx:end_idx]
def __len__(self):
return len(self.data)
def add_noise(self,y):
newimg = random_shapes((256, 256),min_shapes=30,max_shapes=42,
multichannel=False, min_size=20,max_size=30,allow_overlap=True)[0]/255.0
x = y.copy()
x[np.where(newimg < 0.9)] = x[np.where(newimg < 0.9)]+1-newimg[newimg < 0.9]
return x
def __getitem__(self, idx):
y = self.data[idx]
x = torch.FloatTensor(self.add_noise(y)).unsqueeze(0)
return x, torch.FloatTensor(y).unsqueeze(0)
class Denoise(nn.Module):
def __init__(self):
super(Denoise, self).__init__()
# encoder layers
self.enc1 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
self.enc2 = nn.Conv2d(256, 32, kernel_size=3, padding=1)
self.enc3 = nn.Conv2d(32, 16, kernel_size=3, padding=1)
self.enc4 = nn.Conv2d(16, 8, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(2, 2)
# decoder layers
self.dec1 = nn.ConvTranspose2d(8, 8, kernel_size=3, stride=1)
self.dec2 = nn.ConvTranspose2d(8, 16, kernel_size=3, stride=1)
self.dec3 = nn.ConvTranspose2d(16, 32, kernel_size=3, stride=1)
self.dec4 = nn.ConvTranspose2d(32, 256, kernel_size=3, stride=1)
self.out = nn.Conv2d(256, 256, kernel_size=3, padding=1)
def forward(self, x):
# encode
x = F.relu(self.enc1(x))
x = self.pool(x)
x = F.relu(self.enc2(x))
x = self.pool(x)
x = F.relu(self.enc3(x))
x = self.pool(x)
x = F.relu(self.enc4(x))
x = self.pool(x) # the latent space representation
# decode
x = F.relu(self.dec1(x))
x = F.relu(self.dec2(x))
x = F.relu(self.dec3(x))
x = F.relu(self.dec4(x))
x = F.sigmoid(self.out(x))
return x