115 lines
3.6 KiB
Python
115 lines
3.6 KiB
Python
import os
|
|
import sys
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.nn import functional as F
|
|
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
|
from models.convnext import ConvNeXt, CNBlockConfig
|
|
|
|
class ViewEncoder(nn.Module):
|
|
def __init__(self, features=128, device='cuda'):
|
|
super().__init__()
|
|
self.features: int = features
|
|
|
|
self.flatten = nn.Flatten()
|
|
self.lin1 = nn.LazyLinear(features).to(device)
|
|
self.act1 = nn.GELU().to(device)
|
|
self.lin2 = nn.Linear(features, features).to(device)
|
|
self.act2 = nn.Tanh().to(device)
|
|
|
|
|
|
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
|
x: torch.Tensor = self.flatten(input)
|
|
x= self.lin1(x)
|
|
x = self.act1(x)
|
|
x = self.lin2(x)
|
|
x = self.act2(x)
|
|
x = torch.ones((x.shape[0], self.features, 8, 8), device=x.get_device()) * x[:,:,None,None]
|
|
|
|
return x
|
|
|
|
|
|
from torchvision.ops.misc import Conv2dNormActivation, Permute
|
|
|
|
class DumbReconstructor(nn.Module):
|
|
def __init__(self, view_features=128, device='cuda'):
|
|
super().__init__()
|
|
|
|
# self.model = torch.nn.Sequential(
|
|
# nn.ConvTranspose2d(768 + view_features, 128, 7, 4, 3),
|
|
# nn.BatchNorm2d(128),
|
|
# nn.ConvTranspose2d(128, 128, 7, 4, 3),
|
|
# nn.Conv2d(128, 64, 7, 1, 'same'),
|
|
# nn.BatchNorm2d(64),
|
|
# nn.ConvTranspose2d(64, 32, 7, 4, 3),
|
|
# nn.Conv2d(32, 1, 5, 1, 'same')
|
|
# ).to(device)
|
|
|
|
self.last = nn.Sequential(
|
|
nn.Linear(in_features=1024, out_features=480*640, bias=True),
|
|
).to(device)
|
|
|
|
def forward(self, input) -> torch.Tensor:
|
|
x = input
|
|
x = self.last(x).view((1,1,480,640))
|
|
return x
|
|
|
|
|
|
class DumbEncoder(nn.Module):
|
|
def __init__(self, view_features=128, device='cuda'):
|
|
super().__init__()
|
|
|
|
self.last = nn.Sequential(
|
|
nn.Flatten(),
|
|
nn.Linear(in_features=3*480*640, out_features=1024, bias=True),
|
|
).to(device)
|
|
|
|
def forward(self, input) -> torch.Tensor:
|
|
x = input
|
|
x = self.last(x)
|
|
return x, None
|
|
|
|
|
|
class DepthEstimation(nn.Module):
|
|
def __init__(self, view_features=128, device='cuda', use_skip=False):
|
|
super().__init__()
|
|
|
|
|
|
block_setting = [
|
|
CNBlockConfig(96, 192, 3),
|
|
CNBlockConfig(192, 384, 3),
|
|
CNBlockConfig(384, 768, 9),
|
|
CNBlockConfig(768, None, 3),
|
|
]
|
|
|
|
block_setting_reverse = [
|
|
CNBlockConfig(768 + view_features, 384, 3),
|
|
CNBlockConfig(384 * ( 2 if use_skip else 1), 192, 9),
|
|
CNBlockConfig(192 * ( 2 if use_skip else 1), 96, 3),
|
|
CNBlockConfig(96 * ( 2 if use_skip else 1), 27, 3),
|
|
CNBlockConfig(27, 9, 3),
|
|
CNBlockConfig(9, None, 3),
|
|
]
|
|
|
|
self.convnext = ConvNeXt(block_setting, device=device)
|
|
# self.convnext = DumbEncoder()
|
|
self.flatten = nn.Flatten()
|
|
self.decoder = ConvNeXt(block_setting_reverse, use_skip=use_skip, reversed=True, device=device)
|
|
# self.decoder = DumbReconstructor()
|
|
|
|
self.view_encoder = ViewEncoder(view_features)
|
|
|
|
|
|
def forward(self, input) -> torch.Tensor:
|
|
image, view = input
|
|
x, features = self.convnext(image)
|
|
view_features = self.view_encoder(view)
|
|
# x = self.flatten(x)
|
|
x = self.decoder([torch.concat((x, view_features), axis=1), features])
|
|
x = torch.tanh(x[0]), x[1]
|
|
# x = self.decoder(x)
|
|
|
|
return x
|