Hi guys,
I want to create my own data generator from a csv
file for the segmentation task. The file contains two columns: 1st column contains the location of the image and the 2nd one contains the location of the mask. I have written the following code:
df = pd.read_csv('/path/to/csv/data.csv')
X = list(df['input_img'])
y = list(df['mask_img'])
X_train, X_valid, y_train, y_valid = train_test_split(
X, y, test_size=0.33, random_state=42)
class ToTensor(object):
"""Convert ndarrays in sample to Tensors."""
def __call__(self, img):
img = img.transpose((2, 0, 1))
# return {'image': torch.from_numpy(img),
# }
return torch.from_numpy(img)
class NumbersDataset(Dataset):
def __init__(self, inputs, labels, transform=None):
classes = [0,1]
self.X = inputs
self.y = labels
self.transform = transform
self.c = 2
def __len__(self):
return len(self.X)
def __getitem__(self, idx):
img_train = cv2.imread(self.X[idx])
img_mask = cv2.imread(self.y[idx])
img_train = cv2.resize(img_train, (427,240), interpolation = cv2.INTER_LANCZOS4)
img_mask = cv2.resize(img_mask, (427,240), interpolation = cv2.INTER_LANCZOS4)
img_mask = cv2.cvtColor(img_mask, cv2.COLOR_BGR2GRAY)
bin_mask = np.zeros_like(img_mask)
bin_mask[(img_mask)>0]=1
bin_mask = bin_mask.reshape(240, 427, 1)
if self.transform:
img_train = self.transform(img_train)
bin_mask = self.transform(bin_mask)
return img_train, bin_mask
if __name__ == '__main__':
dataset_train = NumbersDataset(X_train, y_train, transforms.Compose([ToTensor()]))
# dataset_train = NumbersDataset(X_train, y_train)
dataloader_train = DataLoader(dataset_train, batch_size=4, shuffle=True)
# dataset_valid = NumbersDataset(X_valid, y_valid)
dataset_valid = NumbersDataset(X_valid, y_valid, transforms.Compose([ToTensor()]))
dataloader_valid = DataLoader(dataset_valid, batch_size=4, shuffle=True)
datas = DataBunch(train_dl = dataloader_train, valid_dl = dataloader_valid)
# datas.show_batch()
datas.c = 1
learner = unet_learner(datas, models.resnet34)
I get the following error:
Traceback (most recent call last):
File "dataset_test.py", line 104, in <module>
learner = unet_learner(datas, models.resnet34)
File "/home/sarvagya/miniconda3/envs/gr/lib/python3.6/site-packages/fastai/vision/learner.py", line 121, in unet_learner
bottle=bottle), data.device)
File "/home/sarvagya/miniconda3/envs/gr/lib/python3.6/site-packages/fastai/core.py", line 66, in _init
old_init(self, *args,**kwargs)
File "/home/sarvagya/miniconda3/envs/gr/lib/python3.6/site-packages/fastai/vision/models/unet.py", line 43, in __init__
sfs_szs = model_sizes(encoder, size=imsize)
File "/home/sarvagya/miniconda3/envs/gr/lib/python3.6/site-packages/fastai/callbacks/hooks.py", line 113, in model_sizes
x = dummy_eval(m, size)
File "/home/sarvagya/miniconda3/envs/gr/lib/python3.6/site-packages/fastai/callbacks/hooks.py", line 108, in dummy_eval
return m.eval()(dummy_batch(m, size))
File "/home/sarvagya/miniconda3/envs/gr/lib/python3.6/site-packages/fastai/callbacks/hooks.py", line 104, in dummy_batch
return one_param(m).new(1, ch_in, *size).requires_grad_(False).uniform_(-1.,1.)
TypeError: new() argument after * must be an iterable, not builtin_function_or_method
What should I do?