Returns feature embeddings for images.
(self, images, detach=True, provide_patch_shapes=False)
| 89 | return self._embed(data) |
| 90 | |
| 91 | def _embed(self, images, detach=True, provide_patch_shapes=False): |
| 92 | """Returns feature embeddings for images.""" |
| 93 | |
| 94 | def _detach(features): |
| 95 | if detach: |
| 96 | return [x.detach().cpu().numpy() for x in features] |
| 97 | return features |
| 98 | |
| 99 | _ = self.forward_modules["feature_aggregator"].eval() |
| 100 | with torch.no_grad(): |
| 101 | features = self.forward_modules["feature_aggregator"](images) |
| 102 | |
| 103 | features = [features[layer] for layer in self.layers_to_extract_from] |
| 104 | |
| 105 | features = [ |
| 106 | self.patch_maker.patchify(x, return_spatial_info=True) for x in features |
| 107 | ] |
| 108 | patch_shapes = [x[1] for x in features] |
| 109 | features = [x[0] for x in features] |
| 110 | ref_num_patches = patch_shapes[0] |
| 111 | |
| 112 | for i in range(1, len(features)): |
| 113 | _features = features[i] |
| 114 | patch_dims = patch_shapes[i] |
| 115 | |
| 116 | # TODO(pgehler): Add comments |
| 117 | _features = _features.reshape( |
| 118 | _features.shape[0], patch_dims[0], patch_dims[1], *_features.shape[2:] |
| 119 | ) |
| 120 | _features = _features.permute(0, -3, -2, -1, 1, 2) |
| 121 | perm_base_shape = _features.shape |
| 122 | _features = _features.reshape(-1, *_features.shape[-2:]) |
| 123 | _features = F.interpolate( |
| 124 | _features.unsqueeze(1), |
| 125 | size=(ref_num_patches[0], ref_num_patches[1]), |
| 126 | mode="bilinear", |
| 127 | align_corners=False, |
| 128 | ) |
| 129 | _features = _features.squeeze(1) |
| 130 | _features = _features.reshape( |
| 131 | *perm_base_shape[:-2], ref_num_patches[0], ref_num_patches[1] |
| 132 | ) |
| 133 | _features = _features.permute(0, -2, -1, 1, 2, 3) |
| 134 | _features = _features.reshape(len(_features), -1, *_features.shape[-3:]) |
| 135 | features[i] = _features |
| 136 | features = [x.reshape(-1, *x.shape[-3:]) for x in features] |
| 137 | |
| 138 | # As different feature backbones & patching provide differently |
| 139 | # sized features, these are brought into the correct form here. |
| 140 | features = self.forward_modules["preprocessing"](features) |
| 141 | features = self.forward_modules["preadapt_aggregator"](features) |
| 142 | |
| 143 | if provide_patch_shapes: |
| 144 | return _detach(features), patch_shapes |
| 145 | return _detach(features) |
| 146 | |
| 147 | def fit(self, training_data): |
| 148 | """PatchCore training. |
no test coverage detected