wenmin-wu/ds-skills/3d-encoder-2d-decoder-segmentation

import torch
import torch.nn as nn
import torch.nn.functional as F
class HybridSegModel(nn.Module):
    def __init__(self, encoder_3d, encoder_dims, upscale=4):
        super().__init__()
        self.encoder = encoder_3d  # 3D ResNet returning multi-scale features
        self.decoder = FPNDecoder(encoder_dims, upscale)
    def forward(self, x):
        if x.ndim == 4:
            x = x.unsqueeze(1)  # (B, 1, D, H, W)
        feat_maps_3d = self.encoder(x)  # list of (B, C, D', H', W')
        feat_maps_2d = [f.mean(dim=2) for f in feat_maps_3d]  # pool depth
        return self.decoder(feat_maps_2d)
class FPNDecoder(nn.Module):
    def __init__(self, dims, upscale):
        super().__init__()
        self.convs = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(dims[i]+dims[i-1], dims[i-1], 3, 1, 1, bias=False),
                nn.BatchNorm2d(dims[i-1]), nn.ReLU(inplace=True))
            for i in range(1, len(dims))])
        self.logit = nn.Conv2d(dims[0], 1, 1)
        self.up = nn.Upsample(scale_factor=upscale, mode='bilinear')
    def forward(self, features):
        for i in range(len(features)-1, 0, -1):
            up = F.interpolate(features[i], scale_factor=2, mode='bilinear')
            features[i-1] = self.convs[i-1](torch.cat([features[i-1], up], 1))
        return self.up(self.logit(features[0]))