File size: 2,479 Bytes
f53b39e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
# Copyright (C) 2024 Apple Inc. All Rights Reserved.
# Field of View network architecture.
from typing import Optional
import torch
from torch import nn
from torch.nn import functional as F
class FOVNetwork(nn.Module):
"""Field of View estimation network."""
def __init__(
self,
num_features: int,
fov_encoder: Optional[nn.Module] = None,
):
"""Initialize the Field of View estimation block.
Args:
----
num_features: Number of features used.
fov_encoder: Optional encoder to bring additional network capacity.
"""
super().__init__()
# Create FOV head.
fov_head0 = [
nn.Conv2d(
num_features, num_features // 2, kernel_size=3, stride=2, padding=1
), # 128 x 24 x 24
nn.ReLU(True),
]
fov_head = [
nn.Conv2d(
num_features // 2, num_features // 4, kernel_size=3, stride=2, padding=1
), # 64 x 12 x 12
nn.ReLU(True),
nn.Conv2d(
num_features // 4, num_features // 8, kernel_size=3, stride=2, padding=1
), # 32 x 6 x 6
nn.ReLU(True),
nn.Conv2d(num_features // 8, 1, kernel_size=6, stride=1, padding=0),
]
if fov_encoder is not None:
self.encoder = nn.Sequential(
fov_encoder, nn.Linear(fov_encoder.embed_dim, num_features // 2)
)
self.downsample = nn.Sequential(*fov_head0)
else:
fov_head = fov_head0 + fov_head
self.head = nn.Sequential(*fov_head)
def forward(self, x: torch.Tensor, lowres_feature: torch.Tensor) -> torch.Tensor:
"""Forward the fov network.
Args:
----
x (torch.Tensor): Input image.
lowres_feature (torch.Tensor): Low resolution feature.
Returns:
-------
The field of view tensor.
"""
if hasattr(self, "encoder"):
x = F.interpolate(
x,
size=None,
scale_factor=0.25,
mode="bilinear",
align_corners=False,
)
x = self.encoder(x)[:, 1:].permute(0, 2, 1)
lowres_feature = self.downsample(lowres_feature)
x = x.reshape_as(lowres_feature) + lowres_feature
else:
x = lowres_feature
return self.head(x)
|