from torch import nn
from .md2_feat import Feat_encoder
from .resnet_encoder import ResnetEncoder
from .depth_decoder import CD_DispNet



class CD_Depth(nn.Module):
    def __init__(self, num_layers=50):
        super(CD_Depth, self).__init__()

        self.feat_extractor = Feat_encoder(50)

        self.encoder = ResnetEncoder(50, False, 4)
        self.real_fn = ResnetEncoder(50, True, 1)

        self.decoder = CD_DispNet(self.real_fn.num_ch_enc)

 
    def forward(self, x, is_feat=False, src=False):
        feat_x = self.feat_extractor(x, is_feat)
        if is_feat:
            return feat_x
        feat_x = self.encoder(feat_x, normalize=False)
        real_x = self.real_fn(x)

        return self.decoder(feat_x, real_x)


