def __init__(self, root_dir,shape_data,img_transofrm=None,label_transform=None,nameImageFile="images",nameMasksFile="masks"): self.shape_data=shape_data self.root_dir=root_dir self.img_transofrm = img_transofrm self.label_transform = label_transform self.nameImageFile=nameImageFile self.nameMasksFile=nameMasksFile self.files = list() self.files=os.listdir(self.root_dir) ####Augmentation### self.TrChannels=trs.TransformColorChannels() self.TrFlip=trs.TransformFlip() self.colorJitter=trs.TransformsColorJitter(0.65,0.8,0.8,0.09) self.TrNoisy=trs.TransformNoisyNormal((128,128,128),(128,128,128),0.98,0.02) self.TrBlur=trs.TransformBlur(3) ################### print("Found files: ",len(self.files))
def __init__ (self,rootData,shapeData,imgTransforms=None, labelTransforms=None): self.rootData=rootData self.shapeData=shapeData self.imgTransforms=imgTransforms self.labelTransforms=labelTransforms self.filesDataset=list() ####Augmentation### self.TrChannels=trs.TransformColorChannels() self.TrFlip=trs.TransformFlip() self.colorJitter=trs.TransformsColorJitter(0.65,0.8,0.8,0.09) self.TrNoisy=trs.TransformNoisyNormal((128,128,128),(128,128,128),0.9,0.1) self.TrBlur=trs.TransformBlur(5) ################### rootFiles1=os.listdir(rootData) for fileName in rootFiles1: cropped_train=os.path.join(rootData,fileName,"cropped_train") instrument_datasets=os.listdir(cropped_train) for instrument_dataset in instrument_datasets: imagesPath=os.path.join(cropped_train,instrument_dataset,"images") masksPath=os.path.join(cropped_train,instrument_dataset,"binary_masks") images=os.listdir(imagesPath) masks=os.listdir(masksPath) lenImgs=len(images) lenMasks=len(masks) if lenImgs!=lenMasks: strError = "error of the dimension of the list of images (" + str(lenIms) + ") and masks (" + str(lenMsks) + ")" raise IOError(strError) for i in range(lenImgs): img_file = os.path.join(imagesPath,images[i]) msk_file = os.path.join(masksPath,masks[i]) self.filesDataset.append({ "image":img_file, "label":msk_file })