44 lines
1.1 KiB
Python
44 lines
1.1 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 DepthEstimation(nn.Module):
|
|
def __init__(self, device='cuda'):
|
|
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, 384, 3),
|
|
CNBlockConfig(384, 192, 9),
|
|
CNBlockConfig(192, 96, 3),
|
|
CNBlockConfig(96, 27, 3),
|
|
CNBlockConfig(27, 9, 3),
|
|
CNBlockConfig(9, None, 3),
|
|
]
|
|
|
|
self.convnext = ConvNeXt(block_setting)
|
|
self.flatten = nn.Flatten()
|
|
self.decoder = ConvNeXt(block_setting_reverse, reversed=True)
|
|
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
x, features = self.convnext(x)
|
|
# x = self.flatten(x)
|
|
x = self.decoder([x, features])
|
|
|
|
return x
|