Args: source: [str]. Path to the MVTec data folder. classname: [str or None]. Name of MVTec class that should be provided in this dataset. If None, the datasets iterates over all available images. resize: [int
(
self,
source,
classname,
resize=256,
imagesize=224,
split=DatasetSplit.TRAIN,
train_val_split=1.0,
**kwargs,
)
| 39 | """ |
| 40 | |
| 41 | def __init__( |
| 42 | self, |
| 43 | source, |
| 44 | classname, |
| 45 | resize=256, |
| 46 | imagesize=224, |
| 47 | split=DatasetSplit.TRAIN, |
| 48 | train_val_split=1.0, |
| 49 | **kwargs, |
| 50 | ): |
| 51 | """ |
| 52 | Args: |
| 53 | source: [str]. Path to the MVTec data folder. |
| 54 | classname: [str or None]. Name of MVTec class that should be |
| 55 | provided in this dataset. If None, the datasets |
| 56 | iterates over all available images. |
| 57 | resize: [int]. (Square) Size the loaded image initially gets |
| 58 | resized to. |
| 59 | imagesize: [int]. (Square) Size the resized loaded image gets |
| 60 | (center-)cropped to. |
| 61 | split: [enum-option]. Indicates if training or test split of the |
| 62 | data should be used. Has to be an option taken from |
| 63 | DatasetSplit, e.g. mvtec.DatasetSplit.TRAIN. Note that |
| 64 | mvtec.DatasetSplit.TEST will also load mask data. |
| 65 | """ |
| 66 | super().__init__() |
| 67 | self.source = source |
| 68 | self.split = split |
| 69 | self.classnames_to_use = [classname] if classname is not None else _CLASSNAMES |
| 70 | self.train_val_split = train_val_split |
| 71 | |
| 72 | self.imgpaths_per_class, self.data_to_iterate = self.get_image_data() |
| 73 | |
| 74 | self.transform_img = [ |
| 75 | transforms.Resize(resize), |
| 76 | transforms.CenterCrop(imagesize), |
| 77 | transforms.ToTensor(), |
| 78 | transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), |
| 79 | ] |
| 80 | self.transform_img = transforms.Compose(self.transform_img) |
| 81 | |
| 82 | self.transform_mask = [ |
| 83 | transforms.Resize(resize), |
| 84 | transforms.CenterCrop(imagesize), |
| 85 | transforms.ToTensor(), |
| 86 | ] |
| 87 | self.transform_mask = transforms.Compose(self.transform_mask) |
| 88 | |
| 89 | self.imagesize = (3, imagesize, imagesize) |
| 90 | |
| 91 | def __getitem__(self, idx): |
| 92 | classname, anomaly, image_path, mask_path = self.data_to_iterate[idx] |
nothing calls this directly
no test coverage detected