Spaces:
Running
Running
| # Copyright (C) 2022-present Naver Corporation. All rights reserved. | |
| # Licensed under CC BY-NC-SA 4.0 (non-commercial use only). | |
| # -------------------------------------------------------- | |
| # CroCo model for downstream tasks | |
| # -------------------------------------------------------- | |
| import torch | |
| from .croco import CroCoNet | |
| def croco_args_from_ckpt(ckpt): | |
| if 'croco_kwargs' in ckpt: # CroCo v2 released models | |
| return ckpt['croco_kwargs'] | |
| elif 'args' in ckpt and hasattr(ckpt['args'], 'model'): # pretrained using the official code release | |
| s = ckpt['args'].model # eg "CroCoNet(enc_embed_dim=1024, enc_num_heads=16, enc_depth=24)" | |
| assert s.startswith('CroCoNet(') | |
| return eval('dict'+s[len('CroCoNet'):]) # transform it into the string of a dictionary and evaluate it | |
| else: # CroCo v1 released models | |
| return dict() | |
| class CroCoDownstreamMonocularEncoder(CroCoNet): | |
| def __init__(self, | |
| head, | |
| **kwargs): | |
| """ Build network for monocular downstream task, only using the encoder. | |
| It takes an extra argument head, that is called with the features | |
| and a dictionary img_info containing 'width' and 'height' keys | |
| The head is setup with the croconet arguments in this init function | |
| NOTE: It works by *calling super().__init__() but with redefined setters | |
| """ | |
| super(CroCoDownstreamMonocularEncoder, self).__init__(**kwargs) | |
| head.setup(self) | |
| self.head = head | |
| def _set_mask_generator(self, *args, **kwargs): | |
| """ No mask generator """ | |
| return | |
| def _set_mask_token(self, *args, **kwargs): | |
| """ No mask token """ | |
| self.mask_token = None | |
| return | |
| def _set_decoder(self, *args, **kwargs): | |
| """ No decoder """ | |
| return | |
| def _set_prediction_head(self, *args, **kwargs): | |
| """ No 'prediction head' for downstream tasks.""" | |
| return | |
| def forward(self, img): | |
| """ | |
| img if of size batch_size x 3 x h x w | |
| """ | |
| B, C, H, W = img.size() | |
| img_info = {'height': H, 'width': W} | |
| need_all_layers = hasattr(self.head, 'return_all_blocks') and self.head.return_all_blocks | |
| out, _, _ = self._encode_image(img, do_mask=False, return_all_blocks=need_all_layers) | |
| return self.head(out, img_info) | |
| class CroCoDownstreamBinocular(CroCoNet): | |
| def __init__(self, | |
| head, | |
| **kwargs): | |
| """ Build network for binocular downstream task | |
| It takes an extra argument head, that is called with the features | |
| and a dictionary img_info containing 'width' and 'height' keys | |
| The head is setup with the croconet arguments in this init function | |
| """ | |
| super(CroCoDownstreamBinocular, self).__init__(**kwargs) | |
| head.setup(self) | |
| self.head = head | |
| def _set_mask_generator(self, *args, **kwargs): | |
| """ No mask generator """ | |
| return | |
| def _set_mask_token(self, *args, **kwargs): | |
| """ No mask token """ | |
| self.mask_token = None | |
| return | |
| def _set_prediction_head(self, *args, **kwargs): | |
| """ No prediction head for downstream tasks, define your own head """ | |
| return | |
| def encode_image_pairs(self, img1, img2, return_all_blocks=False): | |
| """ run encoder for a pair of images | |
| it is actually ~5% faster to concatenate the images along the batch dimension | |
| than to encode them separately | |
| """ | |
| ## the two commented lines below is the naive version with separate encoding | |
| #out, pos, _ = self._encode_image(img1, do_mask=False, return_all_blocks=return_all_blocks) | |
| #out2, pos2, _ = self._encode_image(img2, do_mask=False, return_all_blocks=False) | |
| ## and now the faster version | |
| out, pos, _ = self._encode_image( torch.cat( (img1,img2), dim=0), do_mask=False, return_all_blocks=return_all_blocks ) | |
| if return_all_blocks: | |
| out,out2 = list(map(list, zip(*[o.chunk(2, dim=0) for o in out]))) | |
| out2 = out2[-1] | |
| else: | |
| out,out2 = out.chunk(2, dim=0) | |
| pos,pos2 = pos.chunk(2, dim=0) | |
| return out, out2, pos, pos2 | |
| def forward(self, img1, img2): | |
| B, C, H, W = img1.size() | |
| img_info = {'height': H, 'width': W} | |
| return_all_blocks = hasattr(self.head, 'return_all_blocks') and self.head.return_all_blocks | |
| out, out2, pos, pos2 = self.encode_image_pairs(img1, img2, return_all_blocks=return_all_blocks) | |
| if return_all_blocks: | |
| decout = self._decoder(out[-1], pos, None, out2, pos2, return_all_blocks=return_all_blocks) | |
| decout = out+decout | |
| else: | |
| decout = self._decoder(out, pos, None, out2, pos2, return_all_blocks=return_all_blocks) | |
| return self.head(decout, img_info) |