| import math |
| import os |
| import types |
| from typing import List, Tuple, Union |
|
|
| import numpy as np |
| |
| import torch |
| import torch.nn.modules.utils as nn_utils |
| from PIL import Image |
| from torch import nn |
| from torchvision import transforms |
| from tqdm import trange |
|
|
|
|
| def extract_and_save_features( |
| input_img_path_list: List[str], |
| saved_feat_path_list: List[str], |
| img_shape: Tuple[int, int] = (640, 960), |
| stride: int = 8, |
| model_type: str = "dino_vitb8", |
| ) -> Union[np.ndarray, None]: |
| """ |
| Extracts DINO features from a list of images and saves them to disk. |
| |
| Args: |
| input_img_path_list (List[str]): List of image file paths. |
| saved_feat_path_list (List[str]): List of paths to save the extracted features to. |
| img_shape (Tuple[int, int], optional): Image shape to resize to. Defaults to (640, 960). |
| stride (int, optional): Stride for the ViT extractor. Defaults to 8. |
| model_type (str, optional): Type of DINO model to use. Defaults to "dino_vitb8". |
| num_cams (int, optional): Number of cameras. Defaults to 3. |
| """ |
| assert len(input_img_path_list) == len( |
| saved_feat_path_list |
| ), "input_img_path_list and saved_feat_path_list must have the same length." |
| img_shape = list(img_shape) |
| extractor = ViTExtractor( |
| model_type=model_type, |
| stride=stride, |
| ) |
| prep = transforms.Compose( |
| [ |
| transforms.ToTensor(), |
| transforms.Resize(img_shape), |
| transforms.Normalize( |
| mean=[0.485, 0.456, 0.406], |
| std=[0.229, 0.224, 0.225], |
| ), |
| ] |
| ) |
| for i in trange( |
| 0, |
| len(input_img_path_list), |
| desc="Extracting features (will skip existing files)", |
| dynamic_ncols=True, |
| ): |
| image_pth = input_img_path_list[i] |
| feat_pth = saved_feat_path_list[i] |
|
|
| file_exists = os.path.exists(feat_pth) |
| if file_exists: |
| continue |
| os.makedirs(os.path.dirname(feat_pth), exist_ok=True) |
| images = [prep(Image.open(image_pth).convert("RGB"))] |
| preproc_image_lst = torch.stack(images, dim=0).to("cuda") |
| with torch.no_grad(): |
| descriptors = extractor.extract_descriptors( |
| preproc_image_lst, |
| [11], |
| "key", |
| include_cls=False, |
| ) |
| descriptors = descriptors.reshape( |
| descriptors.shape[0], extractor.num_patches[0], extractor.num_patches[1], -1 |
| ).squeeze() |
| descriptors = descriptors.cpu().detach().numpy() |
| if not file_exists: |
| np.save(feat_pth, descriptors) |
| del extractor |
| |
| torch.cuda.empty_cache() |
| return {} |
|
|
|
|
| def delete_features( |
| saved_feat_path_list: List[str], |
| verbose: bool = True, |
| ) -> None: |
| """ |
| Deletes features from disk. |
| """ |
| for i in trange( |
| 0, |
| len(saved_feat_path_list), |
| desc="Deleting features (will skip non-existing files)", |
| dynamic_ncols=True, |
| ): |
| feat_pth = saved_feat_path_list[i] |
|
|
| file_exists = os.path.exists(feat_pth) |
| if file_exists: |
| os.remove(feat_pth) |
| if verbose: |
| print(f"Deleted {feat_pth}") |
| else: |
| continue |
|
|
|
|
| class ViTExtractor: |
| |
| """This class facilitates extraction of features, descriptors, and saliency maps from a ViT. |
| |
| We use the following notation in the documentation of the module's methods: |
| B - batch size |
| h - number of heads. usually takes place of the channel dimension in pytorch's convention BxCxHxW |
| p - patch size of the ViT. either 8 or 16. |
| t - number of tokens. equals the number of patches + 1, e.g. HW / p**2 + 1. Where H and W are the height and width |
| of the input image. |
| d - the embedding dimension in the ViT. |
| """ |
|
|
| def __init__( |
| self, |
| model_type: str = "dino_vits8", |
| stride: int = 4, |
| model: nn.Module = None, |
| device: str = "cuda", |
| ): |
| """ |
| :param model_type: A string specifying the type of model to extract from. |
| [dino_vits8 | dino_vits16 | dino_vitb8 | dino_vitb16 | vit_small_patch8_224 | |
| vit_small_patch16_224 | vit_base_patch8_224 | vit_base_patch16_224] |
| :param stride: stride of first convolution layer. small stride -> higher resolution. |
| :param model: Optional parameter. The nn.Module to extract from instead of creating a new one in ViTExtractor. |
| should be compatible with model_type. |
| """ |
| self.model_type = model_type |
| self.device = device |
| if model is not None: |
| self.model = model |
| else: |
| self.model = ViTExtractor.create_model(model_type) |
|
|
| self.model = ViTExtractor.patch_vit_resolution(self.model, stride=stride) |
| self.model.eval() |
| self.model.to(self.device) |
| print(self.model) |
| p = ( |
| self.model.patch_embed.patch_size |
| if isinstance(self.model.patch_embed.patch_size, int) |
| else self.model.patch_embed.patch_size[0] |
| ) |
| self.p = p |
| self.stride = self.model.patch_embed.proj.stride |
|
|
| self.mean = ( |
| (0.485, 0.456, 0.406) if "dino" in self.model_type else (0.5, 0.5, 0.5) |
| ) |
| self.std = ( |
| (0.229, 0.224, 0.225) if "dino" in self.model_type else (0.5, 0.5, 0.5) |
| ) |
|
|
| self._feats = [] |
| self.hook_handlers = [] |
| self.load_size = None |
| self.num_patches = None |
|
|
| @staticmethod |
| def create_model(model_type: str) -> nn.Module: |
| if "dinov2" in model_type: |
| |
| |
| model = torch.hub.load("facebookresearch/dinov2", model_type) |
| elif "dino" in model_type: |
| model = torch.hub.load("facebookresearch/dino:main", model_type) |
| else: |
| raise NotImplementedError( |
| "Only dino and timm models are supported at the moment." |
| ) |
| return model |
|
|
| @staticmethod |
| def _fix_pos_enc(patch_size: int, stride_hw: Tuple[int, int]): |
| """ |
| Creates a method for position encoding interpolation. |
| :param patch_size: patch size of the model. |
| :param stride_hw: A tuple containing the new height and width stride respectively. |
| :return: the interpolation method |
| """ |
|
|
| def interpolate_pos_encoding( |
| self, x: torch.Tensor, w: int, h: int |
| ) -> torch.Tensor: |
| npatch = x.shape[1] - 1 |
| N = self.pos_embed.shape[1] - 1 |
| if npatch == N and w == h: |
| return self.pos_embed |
| class_pos_embed = self.pos_embed[:, 0] |
| patch_pos_embed = self.pos_embed[:, 1:] |
| dim = x.shape[-1] |
| |
| w0 = 1 + (w - patch_size) // stride_hw[1] |
| h0 = 1 + (h - patch_size) // stride_hw[0] |
| assert ( |
| w0 * h0 == npatch |
| ), f"""got wrong grid size for {h}x{w} with patch_size {patch_size} and |
| stride {stride_hw} got {h0}x{w0}={h0 * w0} expecting {npatch}""" |
| |
| |
| w0, h0 = w0 + 0.1, h0 + 0.1 |
| patch_pos_embed = nn.functional.interpolate( |
| patch_pos_embed.reshape( |
| 1, int(math.sqrt(N)), int(math.sqrt(N)), dim |
| ).permute(0, 3, 1, 2), |
| scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)), |
| mode="bicubic", |
| align_corners=False, |
| recompute_scale_factor=False, |
| ) |
| assert ( |
| int(w0) == patch_pos_embed.shape[-2] |
| and int(h0) == patch_pos_embed.shape[-1] |
| ) |
| patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) |
| return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1) |
|
|
| return interpolate_pos_encoding |
|
|
| @staticmethod |
| def patch_vit_resolution(model: nn.Module, stride: int) -> nn.Module: |
| """ |
| change resolution of model output by changing the stride of the patch extraction. |
| :param model: the model to change resolution for. |
| :param stride: the new stride parameter. |
| :return: the adjusted model |
| """ |
| patch_size = model.patch_embed.patch_size |
| if isinstance(patch_size, tuple): |
| patch_size = patch_size[0] |
| if stride == patch_size: |
| return model |
|
|
| stride = nn_utils._pair(stride) |
| assert all( |
| [(patch_size // s_) * s_ == patch_size for s_ in stride] |
| ), f"stride {stride} should divide patch_size {patch_size}" |
|
|
| |
| model.patch_embed.proj.stride = stride |
| |
| model.interpolate_pos_encoding = types.MethodType( |
| ViTExtractor._fix_pos_enc(patch_size, stride), model |
| ) |
| return model |
|
|
| def preprocess( |
| self, |
| image_path, |
| load_size: Union[int, Tuple[int, int]], |
| ) -> Tuple[torch.Tensor, Image.Image]: |
| """ |
| Preprocesses an image before extraction. |
| :param image_path: path to image to be extracted. |
| :param load_size: optional. Size to resize image before the rest of preprocessing. |
| :return: a tuple containing: |
| (1) the preprocessed image as a tensor to insert the model of shape BxCxHxW. |
| (2) the pil image in relevant dimensions |
| """ |
| image = Image.open(image_path).convert("RGB") |
| prep = transforms.Compose( |
| [ |
| transforms.ToTensor(), |
| transforms.Resize(load_size), |
| transforms.Normalize(mean=self.mean, std=self.std), |
| ] |
| ) |
| prep_img = prep(image)[None, ...] |
| return prep_img |
|
|
| def _get_hook(self, facet: str): |
| """ |
| generate a hook method for a specific block and facet. |
| """ |
| if facet in ["attn", "token"]: |
|
|
| def _hook(model, input, output): |
| self._feats.append(output) |
|
|
| return _hook |
|
|
| if facet == "query": |
| facet_idx = 0 |
| elif facet == "key": |
| facet_idx = 1 |
| elif facet == "value": |
| facet_idx = 2 |
| else: |
| raise TypeError(f"{facet} is not a supported facet.") |
|
|
| def _inner_hook(module, input, output): |
| input = input[0] |
| B, N, C = input.shape |
| qkv = ( |
| module.qkv(input) |
| .reshape(B, N, 3, module.num_heads, C // module.num_heads) |
| .permute(2, 0, 3, 1, 4) |
| ) |
| self._feats.append(qkv[facet_idx]) |
|
|
| return _inner_hook |
|
|
| def _register_hooks(self, layers: List[int], facet: str) -> None: |
| """ |
| register hook to extract features. |
| :param layers: layers from which to extract features. |
| :param facet: facet to extract. One of the following options: ['key' | 'query' | 'value' | 'token' | 'attn'] |
| """ |
| for block_idx, block in enumerate(self.model.blocks): |
| if block_idx in layers: |
| if facet == "token": |
| self.hook_handlers.append( |
| block.register_forward_hook(self._get_hook(facet)) |
| ) |
| elif facet == "attn": |
| self.hook_handlers.append( |
| block.attn.attn_drop.register_forward_hook( |
| self._get_hook(facet) |
| ) |
| ) |
| elif facet in ["key", "query", "value"]: |
| self.hook_handlers.append( |
| block.attn.register_forward_hook(self._get_hook(facet)) |
| ) |
| else: |
| raise TypeError(f"{facet} is not a supported facet.") |
|
|
| def _unregister_hooks(self) -> None: |
| """ |
| unregisters the hooks. should be called after feature extraction. |
| """ |
| for handle in self.hook_handlers: |
| handle.remove() |
| self.hook_handlers = [] |
|
|
| def _extract_features( |
| self, batch: torch.Tensor, layers: List[int] = 11, facet: str = "key" |
| ) -> List[torch.Tensor]: |
| """ |
| extract features from the model |
| :param batch: batch to extract features for. Has shape BxCxHxW. |
| :param layers: layer to extract. A number between 0 to 11. |
| :param facet: facet to extract. One of the following options: ['key' | 'query' | 'value' | 'token' | 'attn'] |
| :return : tensor of features. |
| if facet is 'key' | 'query' | 'value' has shape Bxhxtxd |
| if facet is 'attn' has shape Bxhxtxt |
| if facet is 'token' has shape Bxtxd |
| """ |
| B, C, H, W = batch.shape |
| self._feats = [] |
| self._register_hooks(layers, facet) |
| _ = self.model(batch) |
| self._unregister_hooks() |
| self.load_size = (H, W) |
| self.num_patches = ( |
| 1 + (H - self.p) // self.stride[0], |
| 1 + (W - self.p) // self.stride[1], |
| ) |
| return self._feats |
|
|
| def _log_bin(self, x: torch.Tensor, hierarchy: int = 2) -> torch.Tensor: |
| """ |
| create a log-binned descriptor. |
| :param x: tensor of features. Has shape Bxhxtxd. |
| :param hierarchy: how many bin hierarchies to use. |
| """ |
| B = x.shape[0] |
| num_bins = 1 + 8 * hierarchy |
|
|
| bin_x = x.permute(0, 2, 3, 1).flatten(start_dim=-2, end_dim=-1) |
| bin_x = bin_x.permute(0, 2, 1) |
| bin_x = bin_x.reshape( |
| B, bin_x.shape[1], self.num_patches[0], self.num_patches[1] |
| ) |
| |
| sub_desc_dim = bin_x.shape[1] |
|
|
| avg_pools = [] |
| |
| for k in range(0, hierarchy): |
| |
| win_size = 3**k |
| avg_pool = torch.nn.AvgPool2d( |
| win_size, stride=1, padding=win_size // 2, count_include_pad=False |
| ) |
| avg_pools.append(avg_pool(bin_x)) |
|
|
| bin_x = torch.zeros( |
| (B, sub_desc_dim * num_bins, self.num_patches[0], self.num_patches[1]) |
| ).to(self.device) |
| for y in range(self.num_patches[0]): |
| for x in range(self.num_patches[1]): |
| part_idx = 0 |
| |
| for k in range(0, hierarchy): |
| kernel_size = 3**k |
| for i in range(y - kernel_size, y + kernel_size + 1, kernel_size): |
| for j in range( |
| x - kernel_size, x + kernel_size + 1, kernel_size |
| ): |
| if i == y and j == x and k != 0: |
| continue |
| if ( |
| 0 <= i < self.num_patches[0] |
| and 0 <= j < self.num_patches[1] |
| ): |
| bin_x[ |
| :, |
| part_idx |
| * sub_desc_dim : (part_idx + 1) |
| * sub_desc_dim, |
| y, |
| x, |
| ] = avg_pools[k][:, :, i, j] |
| else: |
| temp_i = max(0, min(i, self.num_patches[0] - 1)) |
| temp_j = max(0, min(j, self.num_patches[1] - 1)) |
| bin_x[ |
| :, |
| part_idx |
| * sub_desc_dim : (part_idx + 1) |
| * sub_desc_dim, |
| y, |
| x, |
| ] = avg_pools[k][:, :, temp_i, temp_j] |
| part_idx += 1 |
| bin_x = ( |
| bin_x.flatten(start_dim=-2, end_dim=-1).permute(0, 2, 1).unsqueeze(dim=1) |
| ) |
| |
| return bin_x |
|
|
| def extract_descriptors( |
| self, |
| batch: torch.Tensor, |
| layer: List[int], |
| facet: str = "key", |
| bin: bool = False, |
| include_cls: bool = False, |
| ) -> torch.Tensor: |
| """ |
| extract descriptors from the model |
| :param batch: batch to extract descriptors for. Has shape BxCxHxW. |
| :param layers: layer to extract. A number between 0 to 11. |
| :param facet: facet to extract. One of the following options: ['key' | 'query' | 'value' | 'token'] |
| :param bin: apply log binning to the descriptor. default is False. |
| :return: tensor of descriptors. Bx1xtxd' where d' is the dimension of the descriptors. |
| """ |
| assert facet in [ |
| "key", |
| "query", |
| "value", |
| "token", |
| ], f"""{facet} is not a supported facet for descriptors. |
| choose from ['key' | 'query' | 'value' | 'token'] """ |
| self._extract_features(batch, layer, facet) |
| x = torch.concat(self._feats) |
| |
| |
| if not include_cls: |
| x = x[:, :, 1:, :] |
| else: |
| assert ( |
| not bin |
| ), "bin = True and include_cls = True are not supported together, set one of them False." |
| if not bin: |
| desc = ( |
| x.permute(0, 2, 3, 1).flatten(start_dim=-2, end_dim=-1).unsqueeze(dim=1) |
| ) |
| else: |
| desc = self._log_bin(x) |
| if "reg" in self.model_type: |
| desc = desc[..., 4:, :] |
| return desc |
|
|
| def extract_saliency_maps(self, batch: torch.Tensor) -> torch.Tensor: |
| """ |
| extract saliency maps. The saliency maps are extracted by averaging several attention heads from the last layer |
| in of the CLS token. All values are then normalized to range between 0 and 1. |
| :param batch: batch to extract saliency maps for. Has shape BxCxHxW. |
| :return: a tensor of saliency maps. has shape Bxt-1 |
| """ |
| assert ( |
| self.model_type == "dino_vits8" |
| ), f"saliency maps are supported only for dino_vits model_type." |
| self._extract_features(batch, [11], "attn") |
| head_idxs = [0, 2, 4, 5] |
| curr_feats = self._feats[0] |
| cls_attn_map = curr_feats[:, head_idxs, 0, 1:].mean(dim=1) |
| temp_mins, temp_maxs = cls_attn_map.min(dim=1)[0], cls_attn_map.max(dim=1)[0] |
| cls_attn_maps = (cls_attn_map - temp_mins) / ( |
| temp_maxs - temp_mins |
| ) |
| return cls_attn_maps |
|
|
