| 23 | self.device = device |
| 24 | |
| 25 | def load( |
| 26 | self, |
| 27 | backbone, |
| 28 | layers_to_extract_from, |
| 29 | device, |
| 30 | input_shape, |
| 31 | pretrain_embed_dimension, |
| 32 | target_embed_dimension, |
| 33 | patchsize=3, |
| 34 | patchstride=1, |
| 35 | anomaly_score_num_nn=1, |
| 36 | featuresampler=patchcore.sampler.IdentitySampler(), |
| 37 | nn_method=patchcore.common.FaissNN(False, 4), |
| 38 | **kwargs, |
| 39 | ): |
| 40 | self.backbone = backbone.to(device) |
| 41 | self.layers_to_extract_from = layers_to_extract_from |
| 42 | self.input_shape = input_shape |
| 43 | |
| 44 | self.device = device |
| 45 | self.patch_maker = PatchMaker(patchsize, stride=patchstride) |
| 46 | |
| 47 | self.forward_modules = torch.nn.ModuleDict({}) |
| 48 | |
| 49 | feature_aggregator = patchcore.common.NetworkFeatureAggregator( |
| 50 | self.backbone, self.layers_to_extract_from, self.device |
| 51 | ) |
| 52 | feature_dimensions = feature_aggregator.feature_dimensions(input_shape) |
| 53 | self.forward_modules["feature_aggregator"] = feature_aggregator |
| 54 | |
| 55 | preprocessing = patchcore.common.Preprocessing( |
| 56 | feature_dimensions, pretrain_embed_dimension |
| 57 | ) |
| 58 | self.forward_modules["preprocessing"] = preprocessing |
| 59 | |
| 60 | self.target_embed_dimension = target_embed_dimension |
| 61 | preadapt_aggregator = patchcore.common.Aggregator( |
| 62 | target_dim=target_embed_dimension |
| 63 | ) |
| 64 | |
| 65 | _ = preadapt_aggregator.to(self.device) |
| 66 | |
| 67 | self.forward_modules["preadapt_aggregator"] = preadapt_aggregator |
| 68 | |
| 69 | self.anomaly_scorer = patchcore.common.NearestNeighbourScorer( |
| 70 | n_nearest_neighbours=anomaly_score_num_nn, nn_method=nn_method |
| 71 | ) |
| 72 | |
| 73 | self.anomaly_segmentor = patchcore.common.RescaleSegmentor( |
| 74 | device=self.device, target_size=input_shape[-2:] |
| 75 | ) |
| 76 | |
| 77 | self.featuresampler = featuresampler |
| 78 | |
| 79 | def embed(self, data): |
| 80 | if isinstance(data, torch.utils.data.DataLoader): |