class Dataset(object):
 #----Initializing the class----
    def __init__(self, train_names, masks_dir, img_dir,train, im_size):
        self.train_names=train_names
        self.masks_dir=masks_dir
        self.img_dir=img_dir
        self.train=train 
        self.noise = Variable(torch.zeros(299,299))
        self.data_augm=data_augm
       
        self.im_size=im_size

   
    def __len__(self):
        return len(self.train_names)
    
    #----To get a dataset item----
    def __getitem__(self, idx):

      image = Image.open(self.img_dir + self.train_names[idx])

      #target = 0

      
      if self.train == True:

        try:
          masks = Image.open(self.masks_dir + os.path.splitext(self.train_names[idx])[0]) 
          target = 1
            #----------------------------------------------------------------------- 
        except:
          target = 0
      else:
        target = 0

      image=np.array(image)
      if image.shape[2] == 4:
        #convert the image from RGBA2R
        image = cv2.cvtColor(image, cv2.COLOR_BGR2BGR)
      image=resize(image, (self.im_size , self.im_size),anti_aliasing=True)
      image = transforms.ToTensor()(image)
      image=image.float() 
      target = torch.Tensor(target)

      sample = {'image': image, 'target': target}
      
      return sample