在PyTorch中使用DeepLabv3进行语义分割的迁移学习（上）

"""Author: Manpreet Singh MinhasContact: msminhas at uwaterloo ca"""frompathlibimportPathfromtypingimportAny, Callable, OptionalimportnumpyasnpfromPILimportImagefromtorchvision.datasets.visionimportVisionDatasetclassSegmentationDataset(VisionDataset):
"""A PyTorch dataset for image segmentation task.The dataset is compatible with torchvision transforms.The transforms passed would be applied to both the Images and Masks."""def__init__(self,
root: str,
image_folder: str,
mask_folder: str,
transforms: Optional[Callable] =None,
seed: int=None,
fraction: float=None,
subset: str=None,
image_color_mode: str="rgb",
mask_color_mode: str="grayscale") ->None:
"""Args:root (str): Root directory path.image_folder (str): Name of the folder that contains the images in the root directory.mask_folder (str): Name of the folder that contains the masks in the root directory.transforms (Optional[Callable], optional): A function/transform that takes ina sample and returns a transformed version.E.g, ``transforms.ToTensor`` for images. Defaults to None.seed (int, optional): Specify a seed for the train and test split for reproducible results. Defaults to None.fraction (float, optional): A float value from 0 to 1 which specifies the validation split fraction. Defaults to None.subset (str, optional): 'Train' or 'Test' to select the appropriate set. Defaults to None.image_color_mode (str, optional): 'rgb' or 'grayscale'. Defaults to 'rgb'.mask_color_mode (str, optional): 'rgb' or 'grayscale'. Defaults to 'grayscale'.Raises:OSError: If image folder doesn't exist in root.OSError: If mask folder doesn't exist in root.ValueError: If subset is not either 'Train' or 'Test'ValueError: If image_color_mode and mask_color_mode are either 'rgb' or 'grayscale'"""super().__init__(root, transforms)
image_folder_path=Path(self.root) /image_foldermask_folder_path=Path(self.root) /mask_folderifnotimage_folder_path.exists():
raiseOSError(f"{image_folder_path} does not exist.")
ifnotmask_folder_path.exists():
raiseOSError(f"{mask_folder_path} does not exist.")
ifimage_color_modenotin ["rgb", "grayscale"]:
raiseValueError(
f"{image_color_mode} is an invalid choice. Please enter from rgb grayscale."          )
ifmask_color_modenotin ["rgb", "grayscale"]:
raiseValueError(
f"{mask_color_mode} is an invalid choice. Please enter from rgb grayscale."          )
self.image_color_mode=image_color_modeself.mask_color_mode=mask_color_modeifnotfraction:
self.image_names=sorted(image_folder_path.glob("*"))
self.mask_names=sorted(mask_folder_path.glob("*"))
else:
ifsubsetnotin ["Train", "Test"]:
raise (ValueError(
f"{subset} is not a valid input. Acceptable values are Train and Test."              ))
self.fraction=fractionself.image_list=np.array(sorted(image_folder_path.glob("*")))
self.mask_list=np.array(sorted(mask_folder_path.glob("*")))
ifseed:
np.random.seed(seed)
indices=np.arange(len(self.image_list))
np.random.shuffle(indices)
self.image_list=self.image_list[indices]
self.mask_list=self.mask_list[indices]
ifsubset=="Train":
self.image_names=self.image_list[:int(
np.ceil(len(self.image_list) * (1-self.fraction)))]
self.mask_names=self.mask_list[:int(
np.ceil(len(self.mask_list) * (1-self.fraction)))]
else:
self.image_names=self.image_list[
int(np.ceil(len(self.image_list) * (1-self.fraction))):]
self.mask_names=self.mask_list[
int(np.ceil(len(self.mask_list) * (1-self.fraction))):]
def__len__(self) ->int:
returnlen(self.image_names)
def__getitem__(self, index: int) ->Any:
image_path=self.image_names[index]
mask_path=self.mask_names[index]
withopen(image_path, "rb") asimage_file, open(mask_path,
"rb") asmask_file:
image=Image.open(image_file)
ifself.image_color_mode=="rgb":
image=image.convert("RGB")
elifself.image_color_mode=="grayscale":
image=image.convert("L")
mask=Image.open(mask_file)
ifself.mask_color_mode=="rgb":
mask=mask.convert("RGB")
elifself.mask_color_mode=="grayscale":
mask=mask.convert("L")
sample= {"image": image, "mask": mask}
ifself.transforms:
sample["image"] =self.transforms(sample["image"])
sample["mask"] =self.transforms(sample["mask"])
returnsample

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
batch_sampler=None, num_workers=0, collate_fn=None,
pin_memory=False, drop_last=False, timeout=0,
worker_init_fn=None, *, prefetch_factor=2,
persistent_workers=False)

data_dir--Train------Image---------Image1---------ImageN------Mask---------Mask1---------MaskN--Train------Image---------Image1---------ImageN------Mask---------Mask1---------MaskN

--data_dir------Image---------Image1---------ImageN------Mask---------Mask1---------MaskN