83 lines
2.4 KiB
Python
83 lines
2.4 KiB
Python
# 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)
|