import math from typing import List, Optional import torch import torch.nn as nn from einops.layers.torch import Rearrange class ConvModule(nn.Module): def __init__(self, in_channels: int, out_channels: int, drop_rate: float = 0.0, bias: bool = False, **kwargs): super().__init__() self.conv = nn.Conv2d(in_channels, out_channels, bias=bias, **kwargs) self.bn = nn.BatchNorm2d(out_channels) self.actv = nn.GELU() self.drop = nn.Dropout2d(drop_rate) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.conv(x) x = self.bn(x) x = self.actv(x) x = self.drop(x) return x class DepthwiseSeparableConvModule(nn.Module): def __init__( self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, activate: bool = True, drop_rate: float = 0.0, bias: bool = False, **kwargs, ): super().__init__() # depthwise conv self.depthwise = nn.Conv2d( in_channels, in_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=in_channels, bias=bias, **kwargs, ) # pointwise conv self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False) # mormalization & activation self.bn = nn.BatchNorm2d(out_channels) if activate: self.actv = nn.GELU() if drop_rate > 0.0: self.drop = nn.Dropout2d(drop_rate) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.depthwise(x) x = self.pointwise(x) x = self.bn(x) if hasattr(self, "actv"): x = self.actv(x) if hasattr(self, "drop"): x = self.drop(x) return x class GuidedChannelAttention(nn.Module): def __init__(self, channels: int, temp: float = 0.1, reduction: int = 16, init_state: Optional[List[float]] = None): super(GuidedChannelAttention, self).__init__() self.temp = temp self.avg_pool = nn.AdaptiveAvgPool2d(1) use_bias = init_state is not None self.fc = nn.Sequential( nn.Linear(channels, channels // reduction, bias=False), nn.GELU(), nn.Linear(channels // reduction, channels, bias=use_bias), nn.Sigmoid(), ) # initialize the init state if init_state is not None: self._apply_init_state(init_state, channels) def _apply_init_state(self, init_state: List[float], channels: int): if len(init_state) != channels: raise ValueError(f"init_state length ({len(init_state)}) must match channels ({channels})") # clam value and calculate inverse sigmoid eps = 1e-6 target_probs = torch.tensor(init_state) target_probs = torch.clamp(target_probs, min=eps, max=1 - eps) bias_values = torch.log(target_probs / (1 - target_probs)) # initialize the weights accordingly second_linear = self.fc[2] with torch.no_grad(): nn.init.trunc_normal_(second_linear.weight, mean=0.0, std=0.1) second_linear.bias.copy_(bias_values) def forward(self, x: torch.Tensor, segm_logits: torch.Tensor): # appply temp and create segmentation class map segm_map = (segm_logits / self.temp).softmax(1) # resize segmentation map to match the target segm_map = nn.functional.interpolate(segm_map, size=x.shape[2:], mode="nearest-exact") # squeeze and excitation block b, c, _, _ = segm_map.size() y = self.avg_pool(segm_map).view(b, c) y = self.fc(y).view(b, c, 1, 1) segm_map = segm_map * y.expand_as(segm_map) attn_map = segm_map.sum(1, keepdim=True) return x * attn_map class CBAM(nn.Module): def __init__(self, channels, reduction=16): super().__init__() self.ca = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(channels, channels // reduction, 1, bias=False), nn.GELU(), nn.Conv2d(channels // reduction, channels, 1, bias=False), nn.Sigmoid(), ) # spatial attention self.sa = nn.Sequential(nn.Conv2d(2, 1, kernel_size=7, padding=3, bias=False), nn.Sigmoid()) def forward(self, x): # apply channel attention x = x * self.ca(x) # apply spatial attention max_self = torch.max(x, dim=1, keepdim=True)[0] avg_self = torch.mean(x, dim=1, keepdim=True) spatial_map = torch.cat([max_self, avg_self], dim=1) return x * self.sa(spatial_map) class FiLM(nn.Module): def __init__(self, num_features: int, num_metadata: int, hidden_dim: int = 32): super().__init__() self.num_features = num_features self.controller = nn.Sequential( nn.Linear(num_metadata, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, num_features * 2) ) def forward(self, x: torch.Tensor, metadata: torch.Tensor) -> torch.Tensor: params = self.controller(metadata) gamma, beta = torch.split(params, self.num_features, dim=1) return (x * gamma) + beta class FiLM2D(nn.Module): def __init__(self, num_channels: int, num_metadata: int, hidden_dim: int = 32): super().__init__() self.num_channels = num_channels self.controller = nn.Sequential( nn.Linear(num_metadata, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, num_channels * 2) ) def forward(self, x, metadata): params = self.controller(metadata) gamma, beta = torch.split(params, self.num_channels, dim=1) gamma = gamma.view(-1, self.num_channels, 1, 1) beta = beta.view(-1, self.num_channels, 1, 1) return (x * gamma) + beta class PatchCompressor(nn.Module): def __init__(self, dim: int = 384): super().__init__() self.query = nn.Parameter(torch.randn(1, 1, dim)) self.key_proj = nn.Linear(dim, dim) self.val_proj = nn.Linear(dim, dim) def forward(self, x: torch.Tensor) -> torch.Tensor: k = self.key_proj(x) v = self.val_proj(x) attn = (self.query @ k.transpose(-2, -1)) * (x.shape[-1] ** -0.5) attn = torch.nn.functional.softmax(attn, dim=-1) out = (attn @ v).squeeze(1) return out class BottomUpFusionLayer(nn.Module): def __init__(self, in_channels_list: List[int], out_channels: int): super().__init__() self.lateral_convs = nn.ModuleList([nn.Conv2d(in_c, out_channels, kernel_size=1) for in_c in in_channels_list]) self.output_conv = DepthwiseSeparableConvModule(out_channels, out_channels, kernel_size=3, padding=1) def forward(self, features: List[torch.Tensor]) -> torch.Tensor: if len(features) != len(self.lateral_convs): raise ValueError(f"Expected {len(self.lateral_convs)} feature maps, got {len(features)}") fused_feature = 0 target_h, target_w = features[-1].shape[-2:] for i, feat in enumerate(features): x = self.lateral_convs[i](feat) if x.shape[-2:] != (target_h, target_w): x = torch.nn.functional.interpolate(x, size=(target_h, target_w), mode="nearest") fused_feature = fused_feature + x out = self.output_conv(fused_feature) return out class FeaturePooler(nn.Module): def __init__( self, in_dim: int, hidden_dim: int, out_dim: float, input_size: int, target_size: int = 3, max_kernel: int = 7, max_stride: int = 3, conv_drop_rate: float = 0.0, pool_mode: str = "mean", use_depthwise: bool = True, ): super().__init__() self.in_dim = in_dim self.out_dim = out_dim self.pool_mode = pool_mode # get constructor block_cls = ConvModule if use_depthwise: block_cls = DepthwiseSeparableConvModule self.conv_reducer = nn.Sequential() layers_config = self.__get_layer_configs( input_dim=input_size, target_dim=target_size, max_kernel=max_kernel, max_stride=max_stride ) for idx, (kernel, stride) in enumerate(layers_config): if idx == 0: block = block_cls(in_dim, hidden_dim, kernel_size=kernel, drop_rate=conv_drop_rate, stride=stride) elif idx == len(layers_config) - 1: block = block_cls(hidden_dim, out_dim, kernel_size=kernel, drop_rate=conv_drop_rate, stride=stride) else: block = block_cls(hidden_dim, hidden_dim, kernel_size=kernel, drop_rate=conv_drop_rate, stride=stride) self.conv_reducer.append(block) def __get_layer_configs(self, input_dim, target_dim, max_kernel=7, max_stride=2): if input_dim <= target_dim: return [] layers = [] current_dim = input_dim while current_dim > target_dim: # Calculate what would happen if we use Max Kernel + Max Stride next_dim_aggressive = math.floor((current_dim - max_kernel) / max_stride) + 1 if next_dim_aggressive >= target_dim: layers.append((max_kernel, max_stride)) current_dim = next_dim_aggressive continue # Total pixels we still need to remove total_reduction_needed = current_dim - target_dim max_reduction_per_layer = max_kernel - 1 num_layers_needed = math.ceil(total_reduction_needed / max_reduction_per_layer) # Distribute the reduction evenly across these layers base_reduction = total_reduction_needed // num_layers_needed remainder = total_reduction_needed % num_layers_needed for i in range(num_layers_needed): reduction = base_reduction + (1 if i < remainder else 0) k = reduction + 1 layers.append((k, 1)) current_dim = target_dim break return layers def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.conv_reducer(x) if self.pool_mode == "sum": return x.sum(dim=(2, 3)) elif self.pool_mode == "mean": return x.mean(dim=(2, 3)) else: return x class BioMassHead(nn.Module): def __init__( self, in_dim: int, hidden: int, targets: int, num_extra_targts: int = 0, reduce_feats: Optional[int] = None, dropout: float = 0.0, use_modulation: bool = False, ): super().__init__() self.hidden = nn.Linear(in_dim, hidden, bias=False) self.norm = nn.BatchNorm1d(hidden) self.actv = nn.GELU() self.drop = nn.Dropout(dropout) if num_extra_targts >= 1 and use_modulation == False: self.reducer = nn.Linear(hidden, reduce_feats) self.fc = nn.Linear(reduce_feats + num_extra_targts, targets) elif num_extra_targts >= 1 and use_modulation == True: self.film_layer = FiLM(num_features=hidden, num_metadata=num_extra_targts, hidden_dim=num_extra_targts * 8) self.fc = nn.Linear(hidden, targets) else: self.fc = nn.Linear(hidden, targets) def forward(self, x: torch.Tensor, *extras: List[List[torch.Tensor]]) -> torch.Tensor: # forwrd pass the feature encoder x = self.hidden(x) x = self.norm(x) x = self.actv(x) # forward pass the final heads if hasattr(self, "reducer"): x = self.drop(x) x = self.reducer(x) x = torch.concat([x, *extras], dim=-1) elif hasattr(self, "film_layer"): metadata = torch.concat(extras, dim=-1) x = self.film_layer(x, metadata) x = self.drop(x) else: x = self.drop(x) return self.fc(x) class AbstractedBioMassHead(nn.Module): def __init__( self, in_dim: int, num_patches: int, num_targets: int, out_feats: Optional[None] = None, drop_rate: float = 0.0, conv_drop_rate: float = 0.0, channel_factor: int = 1, create_specialists: bool = False, # metadata processing parametres use_modulation: bool = False, num_extra_feats: int = 0, # attention related parameters use_attention: bool = True, attn_temp: float = 1.0, segm_dim: Optional[int] = None, attn_init: Optional[List[int]] = None, # parameters required for pooling pool_mode: str = "mean", target_grid: int = 1, grid_size: Optional[int] = None, use_depthwise: bool = True, ): super().__init__() assert pool_mode in ("mean", "strided_mean", "spatial") # build attention layers if use_attention: assert segm_dim is not None self.attention_block = GuidedChannelAttention( channels=segm_dim, temp=attn_temp, reduction=2, init_state=attn_init ) # build channel reducer if out_feats is None: out_feats = in_dim // 4 self.reduce_feats = ConvModule( in_channels=in_dim, out_channels=out_feats, drop_rate=conv_drop_rate, kernel_size=1 ) in_dim = out_feats # build feature poolers if pool_mode == "spatial": hidden_dim = max(1, in_dim // channel_factor) out_dim = max(1, in_dim // channel_factor) assert grid_size is not None, "grid size mandatory for pool mode `spatial`" self.feat_pooler = FeaturePooler( in_dim=in_dim, hidden_dim=hidden_dim, out_dim=out_dim, input_size=grid_size, target_size=target_grid, max_kernel=7, use_depthwise=use_depthwise, conv_drop_rate=conv_drop_rate, ) elif pool_mode == "strided_mean": out_dim = in_dim self.feat_pooler = nn.Sequential( ConvModule(in_dim, in_dim, kernel_size=3, stride=2), nn.AdaptiveAvgPool2d(output_size=(1, 1)), nn.Flatten(), ) elif pool_mode == "mean": out_dim = in_dim self.feat_pooler = nn.Sequential(nn.AdaptiveAvgPool2d(output_size=(1, 1)), nn.Flatten()) # build head self.create_specialists = create_specialists in_dim = out_dim * num_patches if self.create_specialists: self.head = nn.ModuleList( BioMassHead( targets=1, in_dim=in_dim, hidden=in_dim // (channel_factor * 4), reduce_feats=in_dim // (channel_factor * 32), num_extra_targts=num_extra_feats, use_modulation=use_modulation, dropout=drop_rate, ) for _ in range(num_targets) ) else: self.head = BioMassHead( targets=num_targets, in_dim=in_dim, hidden=in_dim // (channel_factor * 2), reduce_feats=in_dim // (channel_factor * 16), num_extra_targts=num_extra_feats, use_modulation=use_modulation, dropout=drop_rate, ) def forward( self, feats: List[torch.Tensor], segm_logits: List[torch.Tensor], *extra_feats: List[List[torch.Tensor]] ) -> torch.Tensor: extrated_feats = [] for x, segm_logit in zip(feats, segm_logits): # apply attention if hasattr(self, "attention_block"): x = self.attention_block(x, segm_logit) # apply reducer and pooler x = self.reduce_feats(x) x = self.feat_pooler(x) extrated_feats.append(x) extrated_feats = torch.concat(extrated_feats, dim=1) if self.create_specialists: out_targets = [head(extrated_feats, *extra_feats) for head in self.head] out_targets = torch.concatenate(out_targets, dim=-1) else: out_targets = self.head(extrated_feats, *extra_feats) return out_targets class AbstractExtraHead(nn.Module): def __init__( self, in_dim: int, patch_w: int, patch_h: int, num_targets: int, reduce_hidden: int = 2, num_blocks: int = 3, drop_rate: float = 0.0, ): super().__init__() self.combiner = Rearrange("n (h w) c ph pw -> n c (h ph) (w pw)", h=patch_h, w=patch_w) # create feature poolers current_dim = in_dim kernel_maps = { 0: 5, } self.feat_pooler = nn.Sequential() for idx in range(num_blocks): current_kernel = kernel_maps.get(idx, 1) block_constructor = ConvModule if current_kernel == 1 else DepthwiseSeparableConvModule block = block_constructor( in_channels=current_dim, out_channels=current_dim // reduce_hidden, kernel_size=current_kernel ) self.feat_pooler.append(block) current_dim = current_dim // reduce_hidden self.feat_pooler.extend([nn.AdaptiveAvgPool2d(output_size=(1, 1)), nn.Flatten()]) # create the head self.head = nn.Sequential( nn.Linear(current_dim, current_dim), nn.BatchNorm1d(current_dim), nn.GELU(), nn.Dropout(drop_rate), nn.Linear(current_dim, num_targets), ) def forward(self, x: List[torch.Tensor]) -> torch.Tensor: x = torch.stack(x, dim=1) x = self.combiner(x) x = self.feat_pooler(x) return self.head(x)