Spaces:
Runtime error
Runtime error
| # Copyright (C) 2022-present Naver Corporation. All rights reserved. | |
| # Licensed under CC BY-NC-SA 4.0 (non-commercial use only). | |
| # -------------------------------------------------------- | |
| # Heads for downstream tasks | |
| # -------------------------------------------------------- | |
| """ | |
| A head is a module where the __init__ defines only the head hyperparameters. | |
| A method setup(croconet) takes a CroCoNet and set all layers according to the head and croconet attributes. | |
| The forward takes the features as well as a dictionary img_info containing the keys 'width' and 'height' | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| from .dpt_block import DPTOutputAdapter | |
| class PixelwiseTaskWithDPT(nn.Module): | |
| """DPT module for CroCo. | |
| by default, hooks_idx will be equal to: | |
| * for encoder-only: 4 equally spread layers | |
| * for encoder+decoder: last encoder + 3 equally spread layers of the decoder | |
| """ | |
| def __init__( | |
| self, | |
| *, | |
| hooks_idx=None, | |
| layer_dims=[96, 192, 384, 768], | |
| output_width_ratio=1, | |
| num_channels=1, | |
| postprocess=None, | |
| **kwargs, | |
| ): | |
| super(PixelwiseTaskWithDPT, self).__init__() | |
| self.return_all_blocks = True # backbone needs to return all layers | |
| self.postprocess = postprocess | |
| self.output_width_ratio = output_width_ratio | |
| self.num_channels = num_channels | |
| self.hooks_idx = hooks_idx | |
| self.layer_dims = layer_dims | |
| def setup(self, croconet): | |
| dpt_args = { | |
| "output_width_ratio": self.output_width_ratio, | |
| "num_channels": self.num_channels, | |
| } | |
| if self.hooks_idx is None: | |
| if hasattr(croconet, "dec_blocks"): # encoder + decoder | |
| step = {8: 3, 12: 4, 24: 8}[croconet.dec_depth] | |
| hooks_idx = [ | |
| croconet.dec_depth + croconet.enc_depth - 1 - i * step | |
| for i in range(3, -1, -1) | |
| ] | |
| else: # encoder only | |
| step = croconet.enc_depth // 4 | |
| hooks_idx = [ | |
| croconet.enc_depth - 1 - i * step for i in range(3, -1, -1) | |
| ] | |
| self.hooks_idx = hooks_idx | |
| print( | |
| f" PixelwiseTaskWithDPT: automatically setting hook_idxs={self.hooks_idx}" | |
| ) | |
| dpt_args["hooks"] = self.hooks_idx | |
| dpt_args["layer_dims"] = self.layer_dims | |
| self.dpt = DPTOutputAdapter(**dpt_args) | |
| dim_tokens = [ | |
| croconet.enc_embed_dim | |
| if hook < croconet.enc_depth | |
| else croconet.dec_embed_dim | |
| for hook in self.hooks_idx | |
| ] | |
| dpt_init_args = {"dim_tokens_enc": dim_tokens} | |
| self.dpt.init(**dpt_init_args) | |
| def forward(self, x, img_info): | |
| out = self.dpt(x, image_size=(img_info["height"], img_info["width"])) | |
| if self.postprocess: | |
| out = self.postprocess(out) | |
| return out | |