in dataloader files, it’s showing error.
Here is the code
class CDataset(Dataset):
def __init__(self, dataset, img_dir, df, transform=None):
self.dataset = dataset
self.img_dir = img_dir
self.df=df
self.transform = transforms.Compose(transform) if transform else None
def __len__(self):
return len(self.dataset)
def return_image_mask(self, file):
ds=dcmread(self.img_dir+file)
arr=ds.pixel_array
pat_id=file[:-4]
if self.df[self.df['patientId']==pat_id].iloc[0]['Target']==0:
mask=None
else:
mask=np.zeros(arr.shape,dtype='int')
temp=self.df[self.df['patientId']==pat_id]
for i in range(len(temp)):
x=temp.iloc[i]['x']
y=temp.iloc[i]['y']
w=temp.iloc[i]['width']
h=temp.iloc[i]['height']
cv2.rectangle(mask, (int(x), int(y)), (int(x+w), int(y+h)),255,-1)
return arr,mask
def __getitem__(self, i):
img,mask = self.return_image_mask(self.dataset[i])
if self.transform:
img = self.transform(img)
if mask is not None:
mask = self.transform(mask)
temp = self.df[self.df['patientId'] == self.dataset[i][:-4]]
return img,mask,temp.iloc[0].Target
transform = [transforms.ToPILImage(),transforms.Resize(448), transforms.ToTensor()]
train_data=CDataset(train, img_folder,train_df,transform)
test_data=CDataset(test, img_folder,train_df,transform)
train_dl = DataLoader(train_data, batch_size=64)
test_dl = DataLoader(test_data, batch_size=64)