(self, image: torch.Tensor, num_tokens: int)
| 126 | return points |
| 127 | |
| 128 | def forward(self, image: torch.Tensor, num_tokens: int) -> Dict[str, torch.Tensor]: |
| 129 | batch_size, _, img_h, img_w = image.shape |
| 130 | device, dtype = image.device, image.dtype |
| 131 | |
| 132 | aspect_ratio = img_w / img_h |
| 133 | base_h, base_w = int((num_tokens / aspect_ratio) ** 0.5), int((num_tokens * aspect_ratio) ** 0.5) |
| 134 | num_tokens = base_h * base_w |
| 135 | |
| 136 | # Backbones encoding |
| 137 | features, cls_token = self.encoder(image, base_h, base_w, return_class_token=True) |
| 138 | features = [features, None, None, None, None] |
| 139 | |
| 140 | # Concat UVs for aspect ratio input |
| 141 | for level in range(5): |
| 142 | uv = normalized_view_plane_uv(width=base_w * 2 ** level, height=base_h * 2 ** level, aspect_ratio=aspect_ratio, dtype=dtype, device=device) |
| 143 | uv = uv.permute(2, 0, 1).unsqueeze(0).expand(batch_size, -1, -1, -1) |
| 144 | if features[level] is None: |
| 145 | features[level] = uv |
| 146 | else: |
| 147 | features[level] = torch.concat([features[level], uv], dim=1) |
| 148 | |
| 149 | # Shared neck |
| 150 | features = self.neck(features) |
| 151 | |
| 152 | # Heads decoding |
| 153 | points, normal, mask = (getattr(self, head)(features)[-1] if hasattr(self, head) else None for head in ['points_head', 'normal_head', 'mask_head']) |
| 154 | metric_scale = self.scale_head(cls_token) if hasattr(self, 'scale_head') else None |
| 155 | |
| 156 | # Resize |
| 157 | points, normal, mask = (F.interpolate(v, (img_h, img_w), mode='bilinear', align_corners=False, antialias=False) if v is not None else None for v in [points, normal, mask]) |
| 158 | |
| 159 | # Remap output |
| 160 | if points is not None: |
| 161 | points = points.permute(0, 2, 3, 1) |
| 162 | points = self._remap_points(points) # slightly improves the performance in case of very large output values |
| 163 | if normal is not None: |
| 164 | normal = normal.permute(0, 2, 3, 1) |
| 165 | normal = F.normalize(normal, dim=-1) |
| 166 | if mask is not None: |
| 167 | mask = mask.squeeze(1).sigmoid() |
| 168 | if metric_scale is not None: |
| 169 | metric_scale = metric_scale.squeeze(1).exp() |
| 170 | |
| 171 | return_dict = { |
| 172 | 'points': points, |
| 173 | 'normal': normal, |
| 174 | 'mask': mask, |
| 175 | 'metric_scale': metric_scale |
| 176 | } |
| 177 | return_dict = {k: v for k, v in return_dict.items() if v is not None} |
| 178 | |
| 179 | return return_dict |
| 180 | |
| 181 | @torch.inference_mode() |
| 182 | def infer( |
no test coverage detected