import torch import torch.nn as nn from .modules.HF_Decoder import HiF_Decoder from .modules.MiT_Encoder import MixVisionTransformer from .modules.Seg_Head import SegmentationHead class DepthPolyp(nn.Module): def __init__( self, in_channels: int = 3, num_classes: int = 2, # 1 for seg, 1 for depth encoder_name: str = 'b0', decoder_channels: int = 256, activation: str = None, upsampling: int = 4, ): super().__init__() # Encoder configurations encoder_configs = { 'b0': { 'embed_dims': [32, 64, 160, 256], 'num_heads': [1, 2, 5, 8], 'mlp_ratios': [4, 4, 4, 4], 'depths': [2, 2, 2, 2], 'sr_ratios': [8, 4, 2, 1], }, 'b1': { 'embed_dims': [64, 128, 320, 512], 'num_heads': [1, 2, 5, 8], 'mlp_ratios': [4, 4, 4, 4], 'depths': [2, 2, 2, 2], 'sr_ratios': [8, 4, 2, 1], }, 'b2': { 'embed_dims': [64, 128, 320, 512], 'num_heads': [1, 2, 5, 8], 'mlp_ratios': [4, 4, 4, 4], 'depths': [3, 4, 6, 3], 'sr_ratios': [8, 4, 2, 1], }, 'b3': { 'embed_dims': [64, 128, 320, 512], 'num_heads': [1, 2, 5, 8], 'mlp_ratios': [4, 4, 4, 4], 'depths': [3, 4, 18, 3], 'sr_ratios': [8, 4, 2, 1], }, 'b4': { 'embed_dims': [64, 128, 320, 512], 'num_heads': [1, 2, 5, 8], 'mlp_ratios': [4, 4, 4, 4], 'depths': [3, 8, 27, 3], 'sr_ratios': [8, 4, 2, 1], }, 'b5': { 'embed_dims': [64, 128, 320, 512], 'num_heads': [1, 2, 5, 8], 'mlp_ratios': [4, 4, 4, 4], 'depths': [3, 6, 40, 3], 'sr_ratios': [8, 4, 2, 1], }, } if encoder_name not in encoder_configs: raise ValueError(f"encoder_name should be one of {list(encoder_configs.keys())}, got {encoder_name}") config = encoder_configs[encoder_name] # Build encoder self.encoder = MixVisionTransformer( in_chans=in_channels, embed_dims=config['embed_dims'], num_heads=config['num_heads'], mlp_ratios=config['mlp_ratios'], qkv_bias=True, depths=config['depths'], sr_ratios=config['sr_ratios'], drop_rate=0.0, drop_path_rate=0.1, ) self.decoder = HiF_Decoder( encoder_channels=config['embed_dims'], decoder_channels=decoder_channels, ) # Build segmentation head (nn.Sequential style) self.segmentation_head = SegmentationHead( in_channels=decoder_channels//4, out_channels=num_classes, activation=activation, kernel_size=1, upsampling=upsampling, ) self.name = f"DepthPolyp-{encoder_name}" def forward(self, x): """Forward pass Args: x: Input tensor of shape (B, C, H, W) Returns: Output tensor of shape (B, num_classes, H, W) """ # Encoder - returns features at [H/4, H/8, H/16, H/32] encoder_features = self.encoder(x) # Decoder - returns features at H/4 fpn_features = self.decoder(encoder_features) decoder_output = fpn_features # print(f"Decoder output shape: {decoder_output.shape}") # Segmentation head - upsample to original size masks = self.segmentation_head(decoder_output) pred_seg = torch.sigmoid(masks[:, 0:1, :, :]) # segmentation 通道 pred_depth = torch.sigmoid(masks[:, 1:2, :, :]) # depth 通道，通常是回归，不做激活 return pred_seg, pred_depth @torch.no_grad() def predict(self, x): """Inference method""" if self.training: self.eval() return self(x) def load_pretrained(self, checkpoint_path, strict=True): """Load pretrained weights Args: checkpoint_path: Path to checkpoint file strict: Whether to strictly enforce key matching """ state_dict = torch.load(checkpoint_path, map_location='cpu') # Handle different checkpoint formats if 'state_dict' in state_dict: state_dict = state_dict['state_dict'] elif 'model' in state_dict: state_dict = state_dict['model'] # Remove module. prefix if present (from DataParallel) new_state_dict = {} for k, v in state_dict.items(): if k.startswith('module.'): new_state_dict[k[7:]] = v else: new_state_dict[k] = v self.load_state_dict(new_state_dict, strict=strict) print(f"✓ Loaded pretrained weights from {checkpoint_path}") def build_depthpolyp( encoder_name='b0', in_channels=3, num_classes=2, decoder_channels=256, activation=None, ): """ Create a DepthPolyp model Args: encoder_name: Encoder variant ('b0', 'b1', 'b2', 'b3', 'b4', 'b5') in_channels: Number of input channels num_classes: Number of output classes decoder_channels: Number of channels in decoder activation: Output activation ('sigmoid', 'softmax', or None) Returns: DepthPolyp model Example: >>> model = build_depthpolyp('b2', num_classes=21, activation='softmax') >>> print(model) """ model = DepthPolyp( in_channels=in_channels, num_classes=num_classes, encoder_name=encoder_name, decoder_channels=decoder_channels, activation=activation, ) return model if __name__ == '__main__': print("="*60) print("Loading Model .....") model = build_depthpolyp( encoder_name='b0', in_channels=3, # Input channels num_classes=2, # Total 2. 1 for seg, 1 for depth decoder_channels=256, activation='sigmoid', ) print("="*60) print("Validating Model .....") print("Check the Param and Complexity(GMACs)") import ptflops macs, params = ptflops.get_model_complexity_info( model, (3, 224, 224), as_strings=True, print_per_layer_stat=False, verbose=False ) print(f" MACs: {macs}, Params: {params}") # output is MACs: 862.17 MMac, Params: 3.57 M print("="*60) print("Check the output .....") dummy_input = torch.randn(1, 3, 224, 224) # B, C, H, W, single RGB image output_seg, output_depth = model(dummy_input) print("input_shape is:", dummy_input.shape) print("output_seg shape is:", output_seg.shape) print("output_depth shape is:", output_depth.shape)