bs = 2
num_epochs = 100
learning_rate = 1e-3
mom = 0.9
class SEMDataTrain(data.Dataset):
def __init__(self, image_path, mask_path, in_size=512, out_size=512 ,transforms = None):
"""
Args:
image_path (str): the path where the image is located
mask_path (str): the path where the mask is located
option (str): decide which dataset to import
"""
# all file names
self.mask_path = mask_path
self.image_path = image_path
self.mask_arr = os.listdir(mask_path)
self.image_arr = os.listdir(image_path)
self.in_size, self.out_size = in_size, out_size
# Calculate len
self.data_len = len(self.mask_arr)
def transform(self, img_as_img, msk_as_img):
# Resize
resize = transforms.Resize(size=(256, 256))
msk_as_img = resize(img_as_img)
msk_as_img = resize(msk_as_img)
def __getitem__(self,idx):
"""Get specific data corresponding to the index
Args:
index (int): index of the data
Returns:
Tensor: specific data on index which is converted to Tensor
"""
"""
# GET IMAGE
"""
single_image_name = self.image_arr[idx]
img_as_img = Image.open(os.path.join(self.image_path,single_image_name))
# img_as_img.show()
img_as_np = np.asarray(img_as_img)
mg_as_img = Image.fromarray(img_as_np)
img_as_img = img_as_img.convert('L')
plt.imshow(img_as_img)
plt.show()
# Augmentation
# flip {0: vertical, 1: horizontal, 2: both, 3: none}
flip_num = randint(0, 3)
img_as_np = flip(img_as_np, flip_num)
# Noise Determine {0: Gaussian_noise, 1: uniform_noise
if randint(0, 1):
# Gaussian_noise
gaus_sd, gaus_mean = randint(0, 20), 0
img_as_np = add_gaussian_noise(img_as_np, gaus_mean, gaus_sd)
else:
# uniform_noise
l_bound, u_bound = randint(-20, 0), randint(0, 20)
img_as_np = add_uniform_noise(img_as_np, l_bound, u_bound)
# Brightness
pix_add = randint(-20, 20)
img_as_np = change_brightness(img_as_np, pix_add)
# Elastic distort {0: distort, 1:no distort}
sigma = randint(6, 12)
# sigma = 4, alpha = 34
img_as_np, seed = add_elastic_transform(img_as_np, alpha=34, sigma=sigma, pad_size=20)
# Crop the image
img_height, img_width = img_as_np.shape[0], img_as_np.shape[1]
pad_size = int((self.in_size - self.out_size)/2)
img_as_np = np.pad(img_as_np, pad_size, mode="symmetric")
y_loc, x_loc = randint(0, img_height-self.out_size), randint(0, img_width-self.out_size)
img_as_np = cropping(img_as_np, crop_size=self.in_size, dim1=y_loc, dim2=x_loc)
'''
# Sanity Check for image
img1 = Image.fromarray(img_as_np)
img1.show()
'''
# Normalize the image
img_as_np = normalization2(img_as_np, max=1, min=0)
img_as_np = np.expand_dims(img_as_np, axis=0) # add additional dimension
img_as_tensor = torch.from_numpy(img_as_np).float() # Convert numpy array to tensor
"""
# GET MASK
"""
single_mask_name = self.mask_arr[idx]
msk_as_img = Image.open(os.path.join(self.mask_path,single_mask_name))
# msk_as_img.show()
msk_as_np = np.asarray(msk_as_img)
# flip the mask with respect to image
msk_as_np = flip(msk_as_np, flip_num)
# elastic_transform of mask with respect to image
# sigma = 4, alpha = 34, seed = from image transformation
msk_as_np, _ = add_elastic_transform(
msk_as_np, alpha=34, sigma=sigma, seed=seed, pad_size=20)
msk_as_np = approximate_image(msk_as_np)
# Crop the mask
msk_as_np = cropping(msk_as_np, crop_size=self.out_size, dim1=y_loc, dim2=x_loc)
'''
Sanity Check for mask
img2 = Image.fromarray(msk_as_np)
img2.show()
'''
# Normalize mask to only 0 and 1
msk_as_np = msk_as_np/255
# msk_as_np = np.expand_dims(msk_as_np, axis=0) # add additional dimension
msk_as_tensor = torch.from_numpy(msk_as_np).long() # Convert numpy array to tensor
return (img_as_tensor, msk_as_tensor)
def __len__(self):
"""
Returns:
length (int): length of the data
"""
return self.data_len
yes it works if with RGB images but sicnce the labels are gray scale it is not work