# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. # -------------------------------------------------------- # References: # timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm # DeiT: https://github.com/facebookresearch/deit # -------------------------------------------------------- import numpy as np import torch import torch.nn as nn from einops import rearrange from mmcv.runner import load_checkpoint,load_state_dict from mmseg.models.builder import BACKBONES, NECKS from timm.models.layers import to_2tuple from timm.models.vision_transformer import Block from typing import List from typing import Optional, Union, Dict from transformers.models.videomae import VideoMAEModel, VideoMAEConfig from transformers.models.videomae.modeling_videomae import get_sinusoid_encoding_table import collections.abc from mmseg.models.builder import LOSSES import mmseg.models.losses.cross_entropy_loss as CE from mmseg.models.losses.utils import get_class_weight import warnings import torch.nn.functional as F from mmseg.models.builder import HEADS from mmseg.models.decode_heads.decode_head import BaseDecodeHead from mmcv.cnn import ConvModule from mmcv.runner import force_fp32 from mmseg.ops import resize from mmseg.models.losses import accuracy import mmseg.models.losses.focal_loss as FL from timm.models.vision_transformer import VisionTransformer as vit from functools import partial from timm.models.layers import trunc_normal_, DropPath from torch import einsum import itertools,math import torch.utils.checkpoint as checkpoint from typing import List, Tuple import logging logger=logging.getLogger(__name__) def _convTranspose2dOutput( input_size: int, stride: int, padding: int, dilation: int, kernel_size: int, output_padding: int, ): """ Calculate the output size of a ConvTranspose2d. Taken from: https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html """ return ( (input_size - 1) * stride - 2 * padding + dilation * (kernel_size - 1) + output_padding + 1 ) def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: torch.Tensor): """ embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D) """ assert embed_dim % 2 == 0 omega = np.arange(embed_dim // 2, dtype=np.float32) omega /= embed_dim / 2.0 omega = 1.0 / 10000**omega # (D/2,) pos = pos.reshape(-1) # (M,) out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product emb_sin = np.sin(out) # (M, D/2) emb_cos = np.cos(out) # (M, D/2) emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) return emb def get_3d_sincos_pos_embed(embed_dim: int, grid_size: tuple, cls_token: bool = False): # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. # -------------------------------------------------------- # Position embedding utils # -------------------------------------------------------- """ grid_size: 3d tuple of grid size: t, h, w return: pos_embed: L, D """ assert embed_dim % 16 == 0 t_size, h_size, w_size = grid_size w_embed_dim = embed_dim // 16 * 6 h_embed_dim = embed_dim // 16 * 6 t_embed_dim = embed_dim // 16 * 4 w_pos_embed = get_1d_sincos_pos_embed_from_grid(w_embed_dim, np.arange(w_size)) h_pos_embed = get_1d_sincos_pos_embed_from_grid(h_embed_dim, np.arange(h_size)) t_pos_embed = get_1d_sincos_pos_embed_from_grid(t_embed_dim, np.arange(t_size)) w_pos_embed = np.tile(w_pos_embed, (t_size * h_size, 1)) h_pos_embed = np.tile(np.repeat(h_pos_embed, w_size, axis=0), (t_size, 1)) t_pos_embed = np.repeat(t_pos_embed, h_size * w_size, axis=0) pos_embed = np.concatenate((w_pos_embed, h_pos_embed, t_pos_embed), axis=1) if cls_token: pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) return pos_embed class PatchEmbed(nn.Module): """Frames of 2D Images to Patch Embedding The 3D version of timm.models.vision_transformer.PatchEmbed """ def __init__( self, img_size: int = 224, patch_size: int = 16, num_frames: int = 3, tubelet_size: int = 1, in_chans: int = 3, embed_dim: int = 768, norm_layer: nn.Module = None, flatten: bool = True, bias: bool = True, ): super().__init__() img_size = to_2tuple(img_size) patch_size = to_2tuple(patch_size) self.img_size = img_size self.patch_size = patch_size self.num_frames = num_frames self.tubelet_size = tubelet_size self.grid_size = ( num_frames // tubelet_size, img_size[0] // patch_size[0], img_size[1] // patch_size[1], ) self.num_patches = self.grid_size[0] * self.grid_size[1] * self.grid_size[2] self.flatten = flatten self.proj = nn.Conv3d( in_chans, embed_dim, kernel_size=(tubelet_size, patch_size[0], patch_size[1]), stride=(tubelet_size, patch_size[0], patch_size[1]), bias=bias, ) self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() def forward(self, x): B, C, T, H, W = x.shape assert ( H == self.img_size[0] ), f"Input image height ({H}) doesn't match model ({self.img_size[0]})." assert ( W == self.img_size[1] ), f"Input image width ({W}) doesn't match model ({self.img_size[1]})." x = self.proj(x) Hp, Wp = x.shape[3], x.shape[4] if self.flatten: x = x.flatten(2).transpose(1, 2) # B,C,T,H,W -> B,C,L -> B,L,C x = self.norm(x) return x, Hp, Wp class Norm2d(nn.Module): def __init__(self, embed_dim: int): super().__init__() self.ln = nn.LayerNorm(embed_dim, eps=1e-6) def forward(self, x): x = x.permute(0, 2, 3, 1) x = self.ln(x) x = x.permute(0, 3, 1, 2).contiguous() return x @NECKS.register_module() class GeospatialNeck(nn.Module): """ Neck that transforms the token-based output of transformer into a single embedding suitable for processing with standard layers. Performs 4 ConvTranspose2d operations on the rearranged input with kernel_size=2 and stride=2 """ def __init__( self, embed_dim: int, first_conv_channels: int, Hp: int = 14, Wp: int = 14, channel_reduction_factor: int = 2, num_convs: int = 4, num_convs_per_upscale: int = 1, dropout: bool = False, drop_cls_token: bool = True, ): """ Args: embed_dim (int): Input embedding dimension first_conv_channel (int): Number of channels for first dimension Hp (int, optional): Height (in patches) of embedding to be upscaled. Defaults to 14. Wp (int, optional): Width (in patches) of embedding to be upscaled. Defaults to 14. channel_reduction_factor (int): Factor that each convolutional block reduces number of channels by. num_convs (int): Number of convolutional upscaling blocks. Each upscales 2x. drop_cls_token (bool, optional): Whether there is a cls_token, which should be dropped. This assumes the cls token is the first token. Defaults to True. """ super().__init__() self.drop_cls_token = drop_cls_token self.Hp = Hp self.Wp = Wp self.H_out = Hp self.W_out = Wp self.dropout = dropout conv_kernel_size = 3 conv_padding = 1 kernel_size = 2 stride = 2 dilation = 1 padding = 0 output_padding = 0 self.embed_dim = embed_dim self.channels = [first_conv_channels // (channel_reduction_factor ** i) for i in range(num_convs)] self.channels = [embed_dim] + self.channels for _ in range(len(self.channels) - 1): self.H_out = _convTranspose2dOutput( self.H_out, stride, padding, dilation, kernel_size, output_padding ) self.W_out = _convTranspose2dOutput( self.W_out, stride, padding, dilation, kernel_size, output_padding ) def _build_upscale_block(channels_in, channels_out): layers = [] layers.append(nn.ConvTranspose2d( channels_in, channels_out, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding, output_padding=output_padding, )) layers += [nn.Sequential( nn.Conv2d(channels_out, channels_out, kernel_size=conv_kernel_size, padding=conv_padding), nn.BatchNorm2d(channels_out), nn.Dropout() if self.dropout else nn.Identity(), nn.ReLU()) for _ in range(num_convs_per_upscale)] return nn.Sequential(*layers) self.layers = nn.ModuleList([ _build_upscale_block(self.channels[i], self.channels[i+1]) for i in range(len(self.channels) - 1) ]) def forward(self, x): x = x[0] if self.drop_cls_token: x = x[:, 1:, :] x = x.permute(0, 2, 1).reshape(x.shape[0], -1, self.Hp, self.Wp) for layer in self.layers: x = layer(x) x = x.reshape((x.shape[0], self.channels[-1], self.H_out, self.W_out)) out = tuple([x]) return out # Added some module class convBlock(nn.Module): def __init__(self,in_channels,out_channels): super(convBlock,self).__init__() self.conv=nn.Sequential( nn.Conv2d(in_channels,out_channels,kernel_size=3,stride=1,padding=1,bias=True), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), nn.Conv2d(out_channels,out_channels,kernel_size=3,stride=1,padding=1,bias=True), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), ) def forward(self,x): x=self.conv(x) return x class upConv(nn.Module): def __init__(self,ch_in,ch_out): super(upConv,self).__init__() self.up=nn.Sequential( nn.Upsample(scale_factor=2), nn.Conv2d(ch_in,ch_out,kernel_size=3,stride=1,padding=1,bias=True), nn.BatchNorm2d(ch_out), nn.ReLU(inplace=True) ) def forward(self,x): x=self.up(x) return x class channelAttn(nn.Module): def __init__(self,channels=18,ratio=6): super(channelAttn,self).__init__() self.avg_pool=nn.AdaptiveAvgPool2d(output_size=1) self.max_pool=nn.AdaptiveMaxPool2d(output_size=1) self.shared_mlp=nn.Sequential( nn.Conv2d(channels,channels//ratio,kernel_size=1,bias=False), nn.ReLU(), nn.Conv2d(channels//ratio,channels,kernel_size=1,bias=False) ) self.sigmoid=nn.Sigmoid() def forward(self,x): avgout=self.shared_mlp(self.avg_pool(x)) maxout=self.shared_mlp(self.max_pool(x)) attn = self.sigmoid(avgout+maxout) out=attn*x return out # Added CNN branch class myNet(nn.Module): def __init__(self,in_channels): super(myNet,self).__init__() self.conv1=convBlock(in_channels,in_channels*2) self.conv2=convBlock(in_channels*2,in_channels*4) self.maxpool=nn.MaxPool2d(kernel_size=2,stride=2) self.up1=upConv(in_channels*4,in_channels*2) self.up_conv1=convBlock(in_channels*2+in_channels*2,in_channels*2) self.last=nn.Conv2d(in_channels*2,in_channels,kernel_size=1,stride=1,padding=0) self.ch_attn1=channelAttn(in_channels*2) def forward(self,x): conv1=self.conv1(x) conv1=self.ch_attn1(conv1)+conv1 x=self.maxpool(conv1) x=self.conv2(x) x=self.up1(x) x=torch.cat([x,conv1],dim=1) x=self.up_conv1(x) out=self.last(x) return out # modify the neck module @NECKS.register_module() class ConvTransformerTokensToEmbeddingNeck(nn.Module): """ Neck that transforms the token-based output of transformer into a single embedding suitable for processing with standard layers. Performs 4 ConvTranspose2d operations on the rearranged input with kernel_size=2 and stride=2 """ def __init__( self, embed_dim: int, output_embed_dim: int, Hp: int = 14, Wp: int = 14, drop_cls_token: bool = True, num_frames:int=3, in_channels:int=6, feature_cat:bool=False, cnn:bool=False, ): """ Args: embed_dim (int): Input embedding dimension output_embed_dim (int): Output embedding dimension Hp (int, optional): Height (in patches) of embedding to be upscaled. Defaults to 14. Wp (int, optional): Width (in patches) of embedding to be upscaled. Defaults to 14. drop_cls_token (bool, optional): Whether there is a cls_token, which should be dropped. This assumes the cls token is the first token. Defaults to True. """ super().__init__() self.drop_cls_token = drop_cls_token self.Hp = Hp self.Wp = Wp self.H_out = Hp self.W_out = Wp self.num_frames=num_frames self.in_channels=in_channels self.feature_cat=feature_cat kernel_size = 2 stride = 2 dilation = 1 padding = 0 output_padding = 0 for _ in range(4): self.H_out = _convTranspose2dOutput( self.H_out, stride, padding, dilation, kernel_size, output_padding ) self.W_out = _convTranspose2dOutput( self.W_out, stride, padding, dilation, kernel_size, output_padding ) self.embed_dim = embed_dim self.output_embed_dim = output_embed_dim self.fpn1 = nn.Sequential( nn.ConvTranspose2d( self.embed_dim, self.output_embed_dim, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding, output_padding=output_padding, ), Norm2d(self.output_embed_dim), nn.GELU(), nn.ConvTranspose2d( self.output_embed_dim, self.output_embed_dim, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding, output_padding=output_padding, ), ) self.fpn2 = nn.Sequential( nn.ConvTranspose2d( self.output_embed_dim, self.output_embed_dim, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding, output_padding=output_padding, ), Norm2d(self.output_embed_dim), nn.GELU(), nn.ConvTranspose2d( self.output_embed_dim, self.output_embed_dim, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding, output_padding=output_padding, ), ) # the added cnn self.high_feat=myNet(num_frames*in_channels) self.cnn=cnn def forward(self, x): x_copy=x[1] x = x[0] if self.drop_cls_token: if self.feature_cat: x = torch.cat([x[:, 1:x.shape[1]//2, :], x[:, x.shape[1]//2+1:, :]], dim=1) elif not self.cnn: x = x[:, 1:, :] elif self.cnn: x = x x = x.permute(0, 2, 1).reshape(x.shape[0], -1, self.Hp, self.Wp) x = self.fpn1(x) x = self.fpn2(x) x = x.reshape((-1, self.output_embed_dim, self.H_out, self.W_out)) #[_,_,224,224] x_copy=torch.reshape(x_copy,(-1,self.num_frames*self.in_channels,224,224)) x_high=self.high_feat(x_copy) x=torch.cat((x,x_high),dim=1) out = tuple([x]) return out @BACKBONES.register_module() class TemporalViTEncoder(nn.Module): """Encoder from an ViT with capability to take in temporal input. This class defines an encoder taken from a ViT architecture. """ def __init__( self, img_size: int = 224, patch_size: int = 16, num_frames: int = 1, tubelet_size: int = 1, in_chans: int = 3, embed_dim: int = 1024, depth: int = 24, num_heads: int = 16, mlp_ratio: float = 4.0, norm_layer: nn.Module = nn.LayerNorm, norm_pix_loss: bool = False, pretrained: str = None ): """ Args: img_size (int, optional): Input image size. Defaults to 224. patch_size (int, optional): Patch size to be used by the transformer. Defaults to 16. num_frames (int, optional): Number of frames (temporal dimension) to be input to the encoder. Defaults to 1. tubelet_size (int, optional): Tubelet size used in patch embedding. Defaults to 1. in_chans (int, optional): Number of input channels. Defaults to 3. embed_dim (int, optional): Embedding dimension. Defaults to 1024. depth (int, optional): Encoder depth. Defaults to 24. num_heads (int, optional): Number of heads used in the encoder blocks. Defaults to 16. mlp_ratio (float, optional): Ratio to be used for the size of the MLP in encoder blocks. Defaults to 4.0. norm_layer (nn.Module, optional): Norm layer to be used. Defaults to nn.LayerNorm. norm_pix_loss (bool, optional): Whether to use Norm Pix Loss. Defaults to False. pretrained (str, optional): Path to pretrained encoder weights. Defaults to None. """ super().__init__() # -------------------------------------------------------------------------- # MAE encoder specifics self.embed_dim = embed_dim self.patch_embed = PatchEmbed( img_size, patch_size, num_frames, tubelet_size, in_chans, embed_dim ) num_patches = self.patch_embed.num_patches self.num_frames = num_frames self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed = nn.Parameter( torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False ) # fixed sin-cos embedding self.blocks = nn.ModuleList( [ Block( embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer, ) for _ in range(depth) ] ) self.norm = norm_layer(embed_dim) self.norm_pix_loss = norm_pix_loss self.pretrained = pretrained self.initialize_weights() def initialize_weights(self): # initialization # initialize (and freeze) pos_embed by sin-cos embedding pos_embed = get_3d_sincos_pos_embed( self.pos_embed.shape[-1], self.patch_embed.grid_size, cls_token=True ) self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0)) # initialize patch_embed like nn.Linear (instead of nn.Conv2d) w = self.patch_embed.proj.weight.data torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1])) if isinstance(self.pretrained, str): self.apply(self._init_weights) print(f"load from {self.pretrained}") load_checkpoint(self, self.pretrained, strict=False, map_location="cpu") elif self.pretrained is None: # # initialize nn.Linear and nn.LayerNorm self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): # we use xavier_uniform following official JAX ViT: torch.nn.init.xavier_uniform_(m.weight) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) def forward(self, x): img=x # embed patches x, _, _ = self.patch_embed(x) # add pos embed w/o cls token x = x + self.pos_embed[:, 1:, :] # append cls token cls_token = self.cls_token + self.pos_embed[:, :1, :] cls_tokens = cls_token.expand(x.shape[0], -1, -1) x = torch.cat((cls_tokens, x), dim=1) # apply Transformer blocks for blk in self.blocks: x = blk(x) x = self.norm(x) # print("geospatial_fm.py x shape:",x.shape) #[b,589,768] return tuple([x,img]) #return feature and original image # defined loss function def class_mean_bce(pred, label, reduction='mean', ignore_index=-100, weight=None, eps=1e-8): if reduction=="mean": op=torch.mean elif reduction=="sum": op=torch.sum() else: raise NotImplementedError target=(label==1).long().unsqueeze(dim=1) anti_target=(label==0).long().unsqueeze(dim=1) xs_pos=torch.sigmoid(pred) xs_neg=1.0-xs_pos # CE loss loss=target*torch.log(xs_pos.clamp(min=eps)) loss.add_(anti_target*torch.log(xs_neg.clamp(min=eps))) pos_loss=torch.div((target*loss).sum(),target.sum()+eps) neg_loss=torch.div((anti_target*loss).sum(),anti_target.sum()+eps) loss=pos_loss+neg_loss return op(-loss) @LOSSES.register_module() class ClassMeanBCE(nn.Module): """reference:https://github.com/moiexpositoalonsolab/deepbiosphere/blob/master/src/deepbiosphere/Losses.py""" def __init__(self, reduction='mean', loss_weight=1.0, ignore_index=-100, eps=1e-8, loss_name="loss_cmbce" ): super(ClassMeanBCE,self).__init__() self.reduction=reduction self.loss_weight=loss_weight self._loss_name=loss_name def forward(self,pred,target,reduction="mean",ignore_index=-100,weight=None,eps=1e-8): loss=self.loss_weight*class_mean_bce( pred, target, reduction=reduction, ignore_index=ignore_index, weight=weight, eps=eps) return loss @property def loss_name(self): """Loss Name. This function must be implemented and will return the name of this loss function. This name will be used to combine different loss items by simple sum operation. In addition, if you want this loss item to be included into the backward graph, `loss_` must be the prefix of the name. Returns: str: The name of this loss item. """ return self._loss_name @LOSSES.register_module() class ClassAwareBCE(CE.CrossEntropyLoss): def __init__(self, use_sigmoid=False, use_mask=False, reduction='mean', class_weight=None, loss_weight=1.0, loss_name='loss_cabce', avg_non_ignore=False): super(ClassAwareBCE, self).__init__() assert (use_sigmoid is False) or (use_mask is False) self.use_sigmoid = use_sigmoid self.use_mask = use_mask self.reduction = reduction self.loss_weight = loss_weight self.class_weight = get_class_weight(class_weight) self.avg_non_ignore = avg_non_ignore if not self.avg_non_ignore and self.reduction == 'mean': warnings.warn( 'Default ``avg_non_ignore`` is False, if you would like to ' 'ignore the certain label and average loss over non-ignore ' 'labels, which is the same with PyTorch official ' 'cross_entropy, set ``avg_non_ignore=True``.') if self.use_sigmoid: self.cls_criterion = CE.binary_cross_entropy elif self.use_mask: self.cls_criterion = CE.mask_cross_entropy else: self.cls_criterion = CE.cross_entropy self._loss_name = loss_name def extra_repr(self): """Extra repr.""" s = f'avg_non_ignore={self.avg_non_ignore}' return s def forward(self, cls_score, label, weight=None, avg_factor=None, reduction_override=None, ignore_index=-100, **kwargs): """Forward function.""" assert reduction_override in (None, 'none', 'mean', 'sum') reduction = ( reduction_override if reduction_override else self.reduction) if self.class_weight is not None: class_weight = cls_score.new_tensor(self.class_weight) else: label_new=label[label!=ignore_index] neg=(1-label_new).sum() pos=label_new.sum() neg_weight=1/(neg/(neg+pos)) pos_weight=1/(pos/(neg+pos)) class_weight=cls_score.new_tensor((neg_weight,pos_weight)) loss_cls = self.loss_weight * self.cls_criterion( cls_score, label, weight, class_weight=class_weight, reduction=reduction, avg_factor=avg_factor, avg_non_ignore=self.avg_non_ignore, ignore_index=ignore_index, **kwargs) return loss_cls @property def loss_name(self): """Loss Name. This function must be implemented and will return the name of this loss function. This name will be used to combine different loss items by simple sum operation. In addition, if you want this loss item to be included into the backward graph, `loss_` must be the prefix of the name. Returns: str: The name of this loss item. """ return self._loss_name @LOSSES.register_module() class EdgeMSELoss(nn.MSELoss): def __init__(self, size_average=None, reduce=None, reduction: str = 'mean', loss_name='loss_mse', loss_weight=1.0) -> None: super(EdgeMSELoss,self).__init__(size_average, reduce, reduction) self._loss_name=loss_name self.loss_weight=loss_weight def forward(self, input, target,weight=None,ignore_index=-100): target=target.unsqueeze(dim=1) # print("mse loss weight",self.loss_weight) loss_mse = self.loss_weight*F.mse_loss(input, target, reduction=self.reduction) return loss_mse @property def loss_name(self): """Loss Name. This function must be implemented and will return the name of this loss function. This name will be used to combine different loss items by simple sum operation. In addition, if you want this loss item to be included into the backward graph, `loss_` must be the prefix of the name. Returns: str: The name of this loss item. """ return self._loss_name # define decoder head @HEADS.register_module() class CropHead(BaseDecodeHead): def __init__(self, num_convs=2, kernel_size=3, concat_input=True, dilation=1, **kwargs): assert num_convs >= 0 and dilation > 0 and isinstance(dilation, int) self.num_convs = num_convs self.concat_input = concat_input self.kernel_size = kernel_size super(CropHead, self).__init__(**kwargs) if num_convs == 0: assert self.in_channels == self.channels conv_padding = (kernel_size // 2) * dilation convs = [] for i in range(num_convs): _in_channels = self.in_channels if i == 0 else self.channels convs.append( ConvModule( _in_channels, self.channels, kernel_size=kernel_size, padding=conv_padding, dilation=dilation, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg)) if len(convs) == 0: self.convs = nn.Identity() else: self.convs = nn.Sequential(*convs) if self.concat_input: self.conv_cat = ConvModule( self.in_channels + self.channels, self.channels, kernel_size=kernel_size, padding=kernel_size // 2, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, act_cfg=self.act_cfg) def _forward_feature(self, inputs): """Forward function for feature maps before classifying each pixel with ``self.cls_seg`` fc. Args: inputs (list[Tensor]): List of multi-level img features. Returns: feats (Tensor): A tensor of shape (batch_size, self.channels, H, W) which is feature map for last layer of decoder head. """ x = self._transform_inputs(inputs) feats = self.convs(x) if self.concat_input: feats = self.conv_cat(torch.cat([x, feats], dim=1)) return feats @force_fp32(apply_to=('seg_logit', )) def losses(self, seg_logit, seg_label): """Compute segmentation loss.""" loss = dict() if self.downsample_label_ratio > 0: seg_label = seg_label.float() target_size = (seg_label.shape[2] // self.downsample_label_ratio, seg_label.shape[3] // self.downsample_label_ratio) seg_label = resize( input=seg_label, size=target_size, mode='nearest') seg_label = seg_label.long() seg_logit = resize( input=seg_logit, size=seg_label.shape[2:], mode='bilinear', align_corners=self.align_corners) if self.sampler is not None: seg_weight = self.sampler.sample(seg_logit, seg_label) else: seg_weight = None seg_label = seg_label.squeeze(1) if not isinstance(self.loss_decode, nn.ModuleList): losses_decode = [self.loss_decode] else: losses_decode = self.loss_decode for loss_decode in losses_decode: if loss_decode.loss_name not in loss: loss[loss_decode.loss_name] = loss_decode( seg_logit, seg_label, weight=seg_weight, ignore_index=self.ignore_index) else: loss[loss_decode.loss_name] += loss_decode( seg_logit, seg_label, weight=seg_weight, ignore_index=self.ignore_index) if len(seg_label.shape)==4: B,C,H,W=seg_label.shape for i in range(C): crop_seg_logit=seg_logit[:,i*2:(i+1)*2,:,:] crop_seg_label=seg_label[:,i,:,:].squeeze() loss["acc_seg_crop_"+str(i+1)]=accuracy( crop_seg_logit,crop_seg_label,ignore_index=self.ignore_index ) elif len(seg_label.shape)==3: crop_seg_logit=seg_logit[:,0:2,:,:] crop_seg_label=seg_label[:,:,:].squeeze() loss["acc_seg_crop"]=accuracy( crop_seg_logit,crop_seg_label,ignore_index=self.ignore_index ) return loss def forward(self, inputs): """Forward function.""" output = self._forward_feature(inputs) output = self.cls_seg(output) return output @LOSSES.register_module() class MultiCropBCE(CE.CrossEntropyLoss): def __init__(self, use_sigmoid=False, use_mask=False, reduction='mean', class_weight=None, loss_weight=1.0, loss_name='loss_multi_crop_bce', avg_non_ignore=False): super(MultiCropBCE, self).__init__() assert (use_sigmoid is False) or (use_mask is False) self.use_sigmoid = use_sigmoid self.use_mask = use_mask self.reduction = reduction self.loss_weight = loss_weight self.class_weight = get_class_weight(class_weight) self.avg_non_ignore = avg_non_ignore if not self.avg_non_ignore and self.reduction == 'mean': print("loss avg_non_ignore:",self.avg_non_ignore) warnings.warn( 'Default ``avg_non_ignore`` is False, if you would like to ' 'ignore the certain label and average loss over non-ignore ' 'labels, which is the same with PyTorch official ' 'cross_entropy, set ``avg_non_ignore=True``.') if self.use_sigmoid: self.cls_criterion = CE.binary_cross_entropy elif self.use_mask: self.cls_criterion = CE.mask_cross_entropy else: self.cls_criterion = CE.cross_entropy self._loss_name = loss_name def extra_repr(self): """Extra repr.""" s = f'avg_non_ignore={self.avg_non_ignore}' return s def forward(self, cls_score, label, weight=None, avg_factor=None, reduction_override=None, ignore_index=-100, **kwargs): """Forward function.""" assert reduction_override in (None, 'none', 'mean', 'sum') reduction = ( reduction_override if reduction_override else self.reduction) if self.class_weight is not None: crops_weight=self.class_weight weight_shape=np.array(crops_weight).shape if len(weight_shape)>1: crop_losses=[] for i in range(weight_shape[0]): crop_cls_score=cls_score[:,i*2:(i+1)*2,:,:] if len(label.shape)==4: crop_label=label[:,i,:,:] elif len(label.shape)==3: crop_label=label crop_weight=crop_cls_score.new_tensor(crops_weight[i]) crop_loss=self.cls_criterion( crop_cls_score, crop_label, weight, class_weight=crop_weight, reduction=reduction, avg_factor=avg_factor, avg_non_ignore=self.avg_non_ignore, ignore_index=ignore_index, **kwargs ) crop_losses.append(crop_loss) loss_cls=sum(crop_losses) return loss_cls else: class_weight=cls_score.new_tensor(self.class_weight) loss_cls = self.loss_weight * self.cls_criterion( cls_score, label, weight, class_weight=class_weight, reduction=reduction, avg_factor=avg_factor, avg_non_ignore=self.avg_non_ignore, ignore_index=ignore_index, **kwargs) else: class_weight=None loss_cls = self.loss_weight * self.cls_criterion( cls_score, label, weight, class_weight=class_weight, reduction=reduction, avg_factor=avg_factor, avg_non_ignore=self.avg_non_ignore, ignore_index=ignore_index, **kwargs) return loss_cls @property def loss_name(self): """Loss Name. This function must be implemented and will return the name of this loss function. This name will be used to combine different loss items by simple sum operation. In addition, if you want this loss item to be included into the backward graph, `loss_` must be the prefix of the name. Returns: str: The name of this loss item. """ return self._loss_name # Others class UnetConvBlock(nn.Module): """ refs:https://github.com/ClarkCGA/multi-temporal-crop-classification-baseline/blob/main/src/models/unet.py This module creates a user-defined number of conv+BN+ReLU layers. Args: in_channels (int)-- number of input features. out_channels (int) -- number of output features. kernel_size (int) -- Size of convolution kernel. stride (int) -- decides how jumpy kernel moves along the spatial dimensions. padding (int) -- how much the input should be padded on the borders with zero. dilation (int) -- dilation ratio for enlarging the receptive field. num_conv_layers (int) -- Number of conv+BN+ReLU layers in the block. drop_rate (float) -- dropout rate at the end of the block. """ def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, num_conv_layers=2, drop_rate=0): super(UnetConvBlock, self).__init__() layers = [nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), ] if num_conv_layers > 1: if drop_rate > 0: layers += [nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), nn.Dropout(drop_rate), ] * (num_conv_layers - 1) else: layers += [nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), ] * (num_conv_layers - 1) self.block = nn.Sequential(*layers) def forward(self, inputs): outputs = self.block(inputs) return outputs class DUC(nn.Module): """ Dense Upscaling Convolution (DUC) layer. Args: in_channels (int): Number of input channels. out_channels (int): Number of output channels. upscale (int): Upscaling factor. Returns: torch.Tensor: Output tensor after applying DUC. """ def __init__(self, in_channels, out_channles, upscale): super(DUC, self).__init__() out_channles = out_channles * (upscale ** 2) self.conv = nn.Conv2d(in_channels, out_channles, 1, bias=False) self.bn = nn.BatchNorm2d(out_channles) self.relu = nn.ReLU(inplace=True) self.pixl_shf = nn.PixelShuffle(upscale_factor=upscale) kernel = self.icnr(self.conv.weight, scale=upscale) self.conv.weight.data.copy_(kernel) def forward(self, x): x = self.relu(self.bn(self.conv(x))) x = self.pixl_shf(x) return x def icnr(self, x, scale=2, init=nn.init.kaiming_normal): """ ICNR (Initialization from Corresponding Normalized Response) function. Args: x (torch.Tensor): Input tensor. scale (int): Upscaling factor. init (function): Initialization function. Returns: torch.Tensor: Initialized kernel. Note: Even with pixel shuffle we still have check board artifacts, the solution is to initialize the d**2 feature maps with the same radom weights: https://arxiv.org/pdf/1707.02937.pdf """ new_shape = [int(x.shape[0] / (scale ** 2))] + list(x.shape[1:]) subkernel = torch.zeros(new_shape) subkernel = init(subkernel) subkernel = subkernel.transpose(0, 1) subkernel = subkernel.contiguous().view(subkernel.shape[0], subkernel.shape[1], -1) kernel = subkernel.repeat(1, 1, scale ** 2) transposed_shape = [x.shape[1]] + [x.shape[0]] + list(x.shape[2:]) kernel = kernel.contiguous().view(transposed_shape) kernel = kernel.transpose(0, 1) return kernel class UnetUpconvBlock(nn.Module): """ Decoder layer decodes the features along the expansive path. Args: in_channels (int) -- number of input features. out_channels (int) -- number of output features. upmode (str) -- Upsampling type. If "fixed" then a linear upsampling with scale factor of two will be applied using bi-linear as interpolation method. If deconv_1 is chosen then a non-overlapping transposed convolution will be applied to upsample the feature maps. If deconv_1 is chosen then an overlapping transposed convolution will be applied to upsample the feature maps. """ def __init__(self, in_channels, out_channels, upmode="deconv_1"): super(UnetUpconvBlock, self).__init__() if upmode == "fixed": layers = [nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True), ] layers += [nn.BatchNorm2d(in_channels), nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False), ] elif upmode == "deconv_1": layers = [nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2, padding=0, dilation=1), ] elif upmode == "deconv_2": layers = [nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, dilation=1), ] # Dense Upscaling Convolution elif upmode == "DUC": up_factor = 2 upsample_dim = (up_factor ** 2) * out_channels layers = [nn.Conv2d(in_channels, upsample_dim, kernel_size=3, padding=1), nn.BatchNorm2d(upsample_dim), nn.ReLU(inplace=True), nn.PixelShuffle(up_factor), ] #layers = [DUC(in_channels, out_channels, upscale=2)] else: raise ValueError("Provided upsampling mode is not recognized.") self.block = nn.Sequential(*layers) def forward(self, inputs): return self.block(inputs) class UnetAdditiveAttentionBlock(nn.Module): r""" additive attention gate (AG) to merge feature maps extracted at multiple scales through skip connection. Args: f_g (int) -- number of feature maps collected from the higher resolution in encoder path. f_x (int) -- number of feature maps in layer "x" in the decoder. f_inter (int) -- number of feature maps after summation equal to the number of learnable multidimensional attention coefficients. Note: Unlike the original paper we upsample """ def __init__(self, F_g, F_x, F_inter): super(UnetAdditiveAttentionBlock, self).__init__() # Decoder self.W_g = nn.Sequential( nn.Conv2d(F_g, F_inter, kernel_size=1, stride=1, padding=0, bias=True), nn.BatchNorm2d(F_inter) ) # Encoder self.W_x = nn.Sequential( nn.Conv2d(F_x, F_inter, kernel_size=1, stride=1, padding=0, bias=True), nn.BatchNorm2d(F_inter) ) # Fused self.psi = nn.Sequential( nn.Conv2d(F_inter, 1, kernel_size=1, stride=1, padding=0, bias=True), nn.BatchNorm2d(1), nn.Sigmoid() ) self.relu = nn.ReLU(inplace=True) def forward(self, g, x): # set_trace() g1 = self.W_g(g) x1 = self.W_x(x) merge = self.relu(g1 + x1) psi = self.psi(merge) return x * psi @BACKBONES.register_module() class Unet(nn.Module): def __init__(self, out_channels, in_channels, filter_config=None, use_skipAtt=False, dropout_rate=0,pretrained=None): """ UNet model with optional additive attention between skip connections for semantic segmentation of multispectral satellite images. Args: n_classes (int): Number of output classes. #我把这个参数改为out_channels,输出通道数，最后通过接个头进行完成分割任务 in_channels (int): Number of input channels. filter_config (tuple, optional): Configuration of filters in the contracting path. Default is None, which uses the configuration (64, 128, 256, 512, 1024, 2048). use_skipAtt (bool, optional): Flag indicating whether to use skip connections with attention. Default is False. dropout_rate (float, optional): Dropout rate applied to the convolutional layers. Default is 0. """ super(Unet, self).__init__() self.in_channels = in_channels self.use_skipAtt = use_skipAtt if not filter_config: filter_config = (64, 128, 256, 512, 1024, 2048) assert len(filter_config) == 6 # Contraction Path self.encoder_1 = UnetConvBlock(self.in_channels, filter_config[0], num_conv_layers=2, drop_rate=dropout_rate) self.encoder_2 = UnetConvBlock(filter_config[0], filter_config[1], num_conv_layers=2, drop_rate=dropout_rate) self.encoder_3 = UnetConvBlock(filter_config[1], filter_config[2], num_conv_layers=2, drop_rate=dropout_rate) self.encoder_4 = UnetConvBlock(filter_config[2], filter_config[3], num_conv_layers=2, drop_rate=dropout_rate) self.encoder_5 = UnetConvBlock(filter_config[3], filter_config[4], num_conv_layers=2, drop_rate=dropout_rate) self.encoder_6 = UnetConvBlock(filter_config[4], filter_config[5], num_conv_layers=2, drop_rate=dropout_rate) self.pool = nn.MaxPool2d(kernel_size=2, stride=2) # Expansion Path self.decoder_1 = UnetUpconvBlock(filter_config[5], filter_config[4], upmode="deconv_2") self.conv1 = UnetConvBlock(filter_config[4] * 2, filter_config[4], num_conv_layers=2, drop_rate=dropout_rate) # self.decoder_2 = UnetUpconvBlock(filter_config[4], filter_config[3], upmode="deconv_2") # self.conv2 = UnetConvBlock(filter_config[3] * 2, filter_config[3], num_conv_layers=2, drop_rate=dropout_rate) # self.decoder_3 = UnetUpconvBlock(filter_config[3], filter_config[2], upmode="deconv_2") # self.conv3 = UnetConvBlock(filter_config[2] * 2, filter_config[2], num_conv_layers=2, drop_rate=dropout_rate) # self.decoder_4 = UnetUpconvBlock(filter_config[2], filter_config[1], upmode="deconv_2") # self.conv4 = UnetConvBlock(filter_config[1] * 2, filter_config[1], num_conv_layers=2, drop_rate=dropout_rate) # self.decoder_5 = UnetUpconvBlock(filter_config[1], filter_config[0], upmode="deconv_2") # self.conv5 = UnetConvBlock(filter_config[0] * 2, filter_config[0], num_conv_layers=2, drop_rate=dropout_rate) # if self.use_skipAtt: # self.Att1 = UnetAdditiveAttentionBlock(F_g=filter_config[4], F_x=filter_config[4], F_inter=filter_config[3]) # self.Att2 = UnetAdditiveAttentionBlock(F_g=filter_config[3], F_x=filter_config[3], F_inter=filter_config[2]) # self.Att3 = UnetAdditiveAttentionBlock(F_g=filter_config[2], F_x=filter_config[2], F_inter=filter_config[1]) # self.Att4 = UnetAdditiveAttentionBlock(F_g=filter_config[1], F_x=filter_config[1], F_inter=filter_config[0]) # self.Att5 = UnetAdditiveAttentionBlock(F_g=filter_config[0], F_x=filter_config[0], # F_inter=int(filter_config[0] / 2)) # self.classifier = nn.Conv2d(filter_config[0], out_channels, kernel_size=1, stride=1, padding=0) # classNumx224x224 self.classifier=nn.Conv2d(filter_config[4], out_channels, kernel_size=1, stride=1, padding=0) self.pretrained = pretrained if pretrained is not None: print(f"load from {self.pretrained}") load_checkpoint(self,self.pretrained,strict=False,map_location="cpu") def forward(self, inputs): """ Forward pass of the UNet model. Args: inputs (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width). Returns: torch.Tensor: Output tensor of shape (batch_size, out_channels, height, width). """ e1 = self.encoder_1(inputs) p1 = self.pool(e1) e2 = self.encoder_2(p1) p2 = self.pool(e2) e3 = self.encoder_3(p2) p3 = self.pool(e3) e4 = self.encoder_4(p3) p4 = self.pool(e4) e5 = self.encoder_5(p4) p5 = self.pool(e5) e6 = self.encoder_6(p5) d6 = self.decoder_1(e6) if self.use_skipAtt: x5 = self.Att1(g=d6, x=e5) skip1 = torch.cat((x5, d6), dim=1) else: skip1 = torch.cat((e5, d6), dim=1) d6_proper = self.conv1(skip1) d5 = self.decoder_2(d6_proper) if self.use_skipAtt: x4 = self.Att2(g=d5, x=e4) skip2 = torch.cat((x4, d5), dim=1) else: skip2 = torch.cat((e4, d5), dim=1) d5_proper = self.conv2(skip2) d4 = self.decoder_3(d5_proper) if self.use_skipAtt: x3 = self.Att3(g=d4, x=e3) skip3 = torch.cat((x3, d4), dim=1) else: skip3 = torch.cat((e3, d4), dim=1) d4_proper = self.conv3(skip3) d3 = self.decoder_4(d4_proper) if self.use_skipAtt: x2 = self.Att4(g=d3, x=e2) skip4 = torch.cat((x2, d3), dim=1) else: skip4 = torch.cat((e2, d3), dim=1) d3_proper = self.conv4(skip4) d2 = self.decoder_5(d3_proper) if self.use_skipAtt: x1 = self.Att5(g=d2, x=e1) skip5 = torch.cat((x1, d2), dim=1) else: skip5 = torch.cat((e1, d2), dim=1) d2_proper = self.conv5(skip5) d1 = self.classifier(d2_proper) return d1