you have helped me integrating , yolo, NAS, optuna, FBN, Transformers in my project which read TIFF SAR imagies and detect vessels and classify it fishing true or false

I use XVIEW3 project
and you have totally help me build new advanced methods as you are a super professional machine learning software engineer

please help me review and deeply integrate the new updates with the original file in the next file without ruining any function or method or the file structure

here's my centernet_unet

from functools import partial
from typing import List, Dict, Any
from omegaconf import DictConfig
from hydra.utils import instantiate, get_class
from pytorch_toolbelt.modules import (
    ACT_RELU,
    decoders as D,
    EncoderModule,
    ResidualDeconvolutionUpsample2d,
)
from torch import nn, Tensor
import torch
import timm
from xview3.centernet.bboxer.functional import centernet_heatmap_nms
import random
import torch.nn.functional as F
from xview3.centernet.models.custom_encoder import CustomTimmEncoder, DebugUNetDecoder
from .unet_blocks import IRBlock, get_unet_block, SqueezeExcitation
from .necks import FPNNeck
from xview3.nas import ProgressiveNAS
from .heads.yolo_head import YOLOHead
from xview3.centernet.models.transformer_classifier import SwinTransformerEncoder
from xview3.centernet.models.necks import FeatureAggregationNetwork
from xview3.centernet.models.heads import DecoupledHeadGroupNormLateShuffle, YOLOHead
from xview3.common import FilterResponseNormalization, ResidualBlock, FeaturePyramidNetwork
from ..constants import (
    CENTERNET_OUTPUT_CLASS_MAP,
    CENTERNET_OUTPUT_OBJECTNESS_MAP,
    CENTERNET_OUTPUT_OFFSET,
    CENTERNET_OUTPUT_SIZE,
    CENTERNET_OUTPUT_VESSEL_MAP,
    CENTERNET_OUTPUT_FISHING_MAP,
    CENTERNET_OUTPUT_YOLO,
    CENTERNET_TARGET_YOLO
)
from xview3.centernet.models.heads.decoupled_head import DecoupledHeadGroupNormLateShuffle
from xview3.centernet.models.heads.yolo_head import YOLOHead
from ..losses.circlenet_loss import CircleNetLoss, YOLOLoss
from xview3.centernet.bboxer.multilabel_circle_coder import MultilabelCircleNetCoder

__all__ = ["CenterNetUNetModel"]
from xview3.utils import TorchDevice
device = TorchDevice.choose_torch_device()

class CenterNetUNetModel(nn.Module):
    def __init__(
        self,
        encoder: EncoderModule,
        decoder_features: List[int],
        heads: List[nn.Module],
        box_coder,
        neck_out_channels: int = 256,
        input_upsample_factor: int = 1,
        activation=ACT_RELU,
        num_extra_blocks=0,
        block_type="UnetBlock",
        upsample_type="bilinear",
        extra_block=partial(IRBlock, stride=2),
        num_classes=2,
        anchors=[(10, 13), (16, 30), (33, 23)],
        nas=None,
    ):
        super().__init__()
        
        # Initial Setup
        self.output_stride = getattr(encoder, 'strides', [32])[0]
        self.initial_upsample = (
            nn.UpsamplingBilinear2d(scale_factor=input_upsample_factor)
            if input_upsample_factor != 1 else nn.Identity()
        )
        unet_block = get_unet_block(block_type, activation=activation)
        self.encoder = encoder
        self.extra_encoder_blocks = nn.ModuleList(
            [extra_block(encoder.channels[-1], encoder.channels[-1]) for _ in range(num_extra_blocks)]
        )
         # Transformer Encoder (Swin Transformer)
        self.encoder = SwinTransformerEncoder(
            img_size=1024,  # Adjust based on your image size
            embed_dim=96,    # Example embed dimension
            depths=[2, 2, 6, 2],
            num_heads=[3, 6, 12, 24],
            window_size=7,
            mlp_ratio=4.,
            qkv_bias=True,
            qk_scale=None,
            drop_rate=0.,
            attn_drop_rate=0.,
            drop_path_rate=0.1,
            norm_layer=nn.LayerNorm,
            in_chans=2  # Adjust based on SAR image channels (e.g., 2 for dual-polarized)
        )

        # Feature Aggregation Network (FAN)
        self.fan = FeatureAggregationNetwork(
            input_channels=[256, 512, 1024, 2048],  # Adjust based on Swin Transformer output stages
            output_channels=256,
            dropout_rate=dropout_rate
        )

        # Decoder
        decoder_layers = []
        current_channels = 256
        for _ in range(num_layers):
            decoder_layers.append(nn.Conv2d(current_channels, num_filters, kernel_size=kernel_size, padding=kernel_size//2))
            if activation == 'relu':
                decoder_layers.append(nn.ReLU())
            else:
                decoder_layers.append(nn.LeakyReLU())
            decoder_layers.append(nn.Dropout2d(p=dropout_rate))
            decoder_layers.append(nn.ConvTranspose2d(num_filters, num_filters, kernel_size=2, stride=2))
            current_channels = num_filters

        self.decoder = nn.Sequential(*decoder_layers)

        # Heads
        self.heads = nn.ModuleList()
        self.heads.append(DecoupledHeadGroupNormLateShuffle(
            classifier_dim=128,
            objectness_dim=128,
            size_dim=128,
            offset_dim=16,
            dropout_rate=dropout_rate,
            activation=activation,
            num_blocks=3
        ))
        self.heads.append(YOLOHead(
            num_classes=num_classes,
            anchors=anchors if anchors else [[10, 13], [16, 30], [33, 23]],
            input_image_size=(1024, 1024),
            stride=32
        ))


        # Upsample Block Selection
        if upsample_type in ['nearest', 'bilinear']:
            upsample_block_cls = partial(
                nn.Upsample,
                scale_factor=2,
                mode=upsample_type,
                align_corners=False
            )
        elif upsample_type == 'rdtsc':
            upsample_block_cls = partial(ResidualDeconvolutionUpsample2d, scale_factor=2)
        elif upsample_type == 'shuffle':
            upsample_block_cls = partial(nn.PixelShuffle, upscale_factor=2)
        else:
            raise ValueError(f"Unknown upsample_type {upsample_type}")

        # Encoder Channels
        if hasattr(encoder, 'channels'):
            encoder_channels = list(encoder.channels)
        else:
            encoder_channels = [32, 24, 40, 112, 320]  # Default channels if not provided
        encoder_channels += [encoder_channels[-1]] * num_extra_blocks

        # Head, NAS, YOLO Head, Classifier Initialization
        self.heads = nn.ModuleList(heads)  # Initialize heads
        self.head_output_channels = getattr(self.heads[0], 'out_channels', 6)  # Adjusted to 6
        self.output_channels = self.head_output_channels  # **Added**

        self.nas = nas  # Initialize nas

        self.yolo_head = YOLOHead(num_classes=num_classes, anchors=anchors)
        self.yolo_output_channels = self.yolo_head.num_anchors * (5 + num_classes)  # Adjusted to 21

        self.classifier = ConvolutionalClassifier(channels=neck_out_channels,
            num_classes=num_classes,
            model_name="resnet50",
            features_only=False)
        self.classifier_output_channels = getattr(self.classifier, 'out_channels', 2)  # Adjusted to 2

        self.box_coder = box_coder

        # Set neck_output_channels directly without multiplication
        self.neck_output_channels = neck_out_channels

        # Decoder Channels Assignment
        self.decoder_channels = decoder_features

        num_encoder_blocks = len(encoder_channels) - 1 

        if len(self.decoder_channels) < num_encoder_blocks:
            self.decoder_channels.extend(
                [self.decoder_channels[-1]] * (num_encoder_blocks - len(self.decoder_channels))
            )
        elif len(self.decoder_channels) > num_encoder_blocks:
            self.decoder_channels = self.decoder_channels[:num_encoder_blocks]

        self.num_decoded_features = len(encoder_channels)

        # Initialize Decoder
        self.decoder = DebugUNetDecoder(
            encoder_channels=encoder_channels,
            decoder_channels=self.decoder_channels,
            unet_block=unet_block,
            upsample_block=upsample_block_cls,
            center=False,
        )

        # Initialize Neck with correct decoder_channels
        self.neck = FPNNeck(channels=self.decoder_channels, out_channels=neck_out_channels)
        # self.neck_output_channels is already set to neck_out_channels

        # Initialize neck_conv with the correct in_channels
        self.neck_conv = nn.Conv2d(
            in_channels=self.neck_output_channels,  # 256
            out_channels=neck_out_channels,         # 256
            kernel_size=1
        )
        self.batch_norm = FilterResponseNormalization(self.neck_output_channels)

        # Calculate total_input_channels based on the actual number of channels
        self.total_input_channels = (
            self.head_output_channels +
            self.yolo_output_channels +
            self.classifier_output_channels
        )
        print(f"Total Input Channels for Fusion: {self.total_input_channels}")
        print(f"head_output_channels: {self.head_output_channels}")
        print(f"yolo_output_channels: {self.yolo_output_channels}")
        print(f"classifier_output_channels: {self.classifier_output_channels}")
        

        # Adjust channels if necessary
        self.yolo_adjust = nn.Conv2d(self.yolo_output_channels, self.head_output_channels, kernel_size=1)
        self.classifier_adjust = nn.Conv2d(self.classifier_output_channels, self.head_output_channels, kernel_size=1)

        # Fusion Layer
        #self.fusion = nn.Conv2d(self.total_input_channels, self.head_output_channels, kernel_size=1)
        self.fusion = nn.Conv2d(18, self.head_output_channels, kernel_size=1)
        self.dropout = nn.Dropout(p=0.5)

        # Initialize Loss Functions
        self.circle_net_loss = CircleNetLoss()
        self.yolo_loss = YOLOLoss(num_classes=num_classes, anchors=anchors)
        self.criterion = YOLOLoss(num_classes, anchors)

    def forward(self, x):
        # Initial Upsampling
        x = self.initial_upsample(x)
        
        # Encoder and Decoder
        feature_maps = self.encoder(x)
        decoded_features = self.decoder(feature_maps)

        # Neck Processing
        neck_features = self.neck(decoded_features)
        neck_features = self.batch_norm(neck_features)
        neck_features = self.dropout(neck_features)
        neck_features = self.neck_conv(neck_features)

        # Heads
        head_outputs = []
        for head in self.heads:
            head_output = head(neck_features)
            head_outputs.append(head_output)

        # Extracting Head Outputs
        objectness_map = head_output[CENTERNET_OUTPUT_OBJECTNESS_MAP]  # [B, 1, H, W]
        vessel_map = head_output[CENTERNET_OUTPUT_VESSEL_MAP]          # [B, 1, H, W]
        fishing_map = head_output[CENTERNET_OUTPUT_FISHING_MAP]        # [B, 1, H, W]
        size_map = head_output[CENTERNET_OUTPUT_SIZE]                  # [B, 1, H, W]
        offset_map = head_output[CENTERNET_OUTPUT_OFFSET]              # [B, 2, H, W]
        yolo_output = head_outputs[1]
        # Adjusting Yolo Output Shape if Necessary
        if len(yolo_output.shape) == 4:
            yolo_output = yolo_output.squeeze(1)  # Remove the extra dimension if it's of size 1

        # Upsample YOLO and Classification Outputs to Match Target Size
        target_size = objectness_map.size()[-2:]  # (H, W)

        if yolo_output.size()[-2:] != target_size:
            yolo_output = F.interpolate(yolo_output, size=target_size, mode='bilinear', align_corners=False)
        
        if classification_output.size()[-2:] != target_size:
            classification_output = F.interpolate(classification_output, size=target_size, mode='bilinear', align_corners=False)

        # Adjust Channels if Necessary
        yolo_output = self.yolo_adjust(yolo_output)               # [B, 6, H, W]
        classification_output = self.classifier_adjust(classification_output)  # [B, 6, H, W]

        # Concatenate All Relevant Feature Maps
        concatenated = torch.cat([
            objectness_map,         # [B, 1, H, W]
            vessel_map,             # [B, 1, H, W]
            fishing_map,            # [B, 1, H, W]
            size_map,               # [B, 1, H, W]
            offset_map,             # [B, 2, H, W]
            yolo_output,            # [B, 6, H, W]
            classification_output   # [B, 6, H, W]
        ], dim=1)  # Total Channels: 1 + 1 + 1 + 1 + 2 + 6 + 6 = 18

        # **Verification Prints:**
        print(f"objectness_map shape: {objectness_map.shape}")          # [B, 1, H, W]
        print(f"vessel_map shape: {vessel_map.shape}")                  # [B, 1, H, W]
        print(f"fishing_map shape: {fishing_map.shape}")                # [B, 1, H, W]
        print(f"size_map shape: {size_map.shape}")                      # [B, 1, H, W]
        print(f"offset_map shape: {offset_map.shape}")                  # [B, 2, H, W]
        print(f"yolo_output shape: {yolo_output.shape}")                # [B, 6, H, W]
        print(f"classification_output shape: {classification_output.shape}")  # [B, 6, H, W]
        print(f"Total concatenated channels: {concatenated.shape[1]}")  # Should be 1+1+1+1+2+6+6 = 18

        # **Assertion for Debugging**
        assert concatenated.shape[1] == 18, \
            f"Expected {self.total_input_channels} channels, but got {concatenated.shape[1]}"

        # Fusion Layer
        fused_output = self.fusion(concatenated)  # [B, 6, H, W]

        # Prepare Outputs
        outputs = {
            CENTERNET_OUTPUT_OBJECTNESS_MAP: objectness_map,
            CENTERNET_OUTPUT_SIZE: size_map,
            CENTERNET_OUTPUT_OFFSET: offset_map,
            CENTERNET_OUTPUT_VESSEL_MAP: vessel_map,
            CENTERNET_OUTPUT_FISHING_MAP: fishing_map,
            CENTERNET_OUTPUT_YOLO: yolo_output,  # Include YOLO output
            # Add `fused_output` if necessary
        }

        return outputs

    def calculate_loss(self, outputs, targets):
        if CENTERNET_TARGET_YOLO not in targets:
            raise KeyError(f"Expected key '{CENTERNET_TARGET_YOLO}' not found in targets.")

        circle_net_loss = self.circle_net_loss(**outputs, **targets)

        if CENTERNET_OUTPUT_YOLO in outputs and CENTERNET_TARGET_YOLO in targets:
            yolo_loss = self.yolo_loss(outputs[CENTERNET_OUTPUT_YOLO], targets[CENTERNET_TARGET_YOLO])
        else:
            yolo_loss = 0

        total_loss = circle_net_loss['objectness'] + circle_net_loss['vessel'] + circle_net_loss['fishing'] + circle_net_loss['size'] + yolo_loss
        return total_loss


    @classmethod
    def from_config(cls, config: DictConfig):
        from hydra.utils import instantiate, get_class
        from omegaconf import OmegaConf
        from omegaconf.errors import ConfigKeyError

        # Enhanced Debugging
        print("Available keys in model config:", config.keys())
        print("Complete model configuration:\n", OmegaConf.to_yaml(config))
        
        # Extract decoder configuration
        decoder_config = config.decoder
        decoder_features = decoder_config.channels  # List of channels, e.g., [256, 384, 512]

        # Check if 'head' exists
        if 'head' not in config:
            raise ConfigKeyError("The key 'head' is missing in the model configuration.")
        
        head_cfg = config.head
        heads = []
        
        if isinstance(head_cfg, list):
            # Handling Multiple Heads
            for single_head_cfg in head_cfg:
                head_type = single_head_cfg.get('type', None) or single_head_cfg.get('_target_', None)
                if head_type is None:
                    raise ConfigKeyError("Both 'type' and '_target_' are missing from the head configuration")

                HeadClass = get_class(head_type)
                
                if head_type == 'xview3.centernet.models.heads.DecoupledHeadGroupNormLateShuffle':
                    channels = single_head_cfg.get('channels', decoder_features[0])  # Use specified channels or fallback
                    classifier_dim = single_head_cfg['classifier_dim']
                    objectness_dim = single_head_cfg['objectness_dim']
                    size_dim = single_head_cfg['size_dim']
                    offset_dim = single_head_cfg['offset_dim']
                    dropout_rate = single_head_cfg.get('dropout_rate', 0.1)
                    activation = single_head_cfg.get('activation', 'relu')  # Default to 'relu' if not specified
                    num_blocks = single_head_cfg['num_blocks']

                    head = HeadClass(
                        channels=channels,
                        classifier_dim=classifier_dim,
                        objectness_dim=objectness_dim,
                        size_dim=size_dim,
                        offset_dim=offset_dim,
                        dropout_rate=dropout_rate,
                        activation=activation,
                        num_blocks=num_blocks,
                    )
                elif head_type == 'xview3.centernet.models.heads.yolo_head.YOLOHead':
                    channels = single_head_cfg.get('channels', decoder_features[0])  # Use specified channels or fallback
                    num_classes = single_head_cfg['num_classes']
                    anchors = single_head_cfg['anchors']

                    head = HeadClass(
                        channels=channels,
                        num_classes=num_classes,
                        anchors=anchors,
                    )
                else:
                    raise NotImplementedError(f"Head type '{head_type}' is not implemented.")

                heads.append(head)
        else:
            # Handling Single Head
            single_head_cfg = head_cfg
            head_type = single_head_cfg.get('type', None) or single_head_cfg.get('_target_', None)
            if head_type is None:
                raise ConfigKeyError("Both 'type' and '_target_' are missing from the head configuration")

            HeadClass = get_class(head_type)

            if head_type == 'xview3.centernet.models.heads.DecoupledHeadGroupNormLateShuffle':
                channels = single_head_cfg.get('channels', decoder_features[0])  # Use specified channels or fallback
                classifier_dim = single_head_cfg['classifier_dim']
                objectness_dim = single_head_cfg['objectness_dim']
                size_dim = single_head_cfg['size_dim']
                offset_dim = single_head_cfg['offset_dim']
                dropout_rate = single_head_cfg.get('dropout_rate', 0.1)
                activation = single_head_cfg.get('activation', 'relu')  # Default to 'relu' if not specified
                num_blocks = single_head_cfg['num_blocks']

                head = HeadClass(
                    channels=channels,
                    classifier_dim=classifier_dim,
                    objectness_dim=objectness_dim,
                    size_dim=size_dim,
                    offset_dim=offset_dim,
                    dropout_rate=dropout_rate,
                    activation=activation,
                    num_blocks=num_blocks,
                )
            elif head_type == 'xview3.centernet.models.heads.yolo_head.YOLOHead':
                channels = single_head_cfg.get('channels', decoder_features[0])  # Use specified channels or fallback
                num_classes = single_head_cfg['num_classes']
                anchors = single_head_cfg['anchors']

                head = HeadClass(
                    channels=channels,
                    num_classes=num_classes,
                    anchors=anchors,
                )
            else:
                raise NotImplementedError(f"Head type '{head_type}' is not implemented.")

            heads.append(head)

        # Debugging: Verify that heads are instantiated
        print(f"[DEBUG] Number of heads instantiated: {len(heads)}")
        for idx, head in enumerate(heads):
            print(f"[DEBUG] Head {idx} type: {type(head)}")

        # Instantiate the encoder
        encoder_config = config.encoder
        encoder: EncoderModule = instantiate(encoder_config)
        
        # Adjust input channels if necessary
        num_channels = config.get('num_channels', 3)
        if num_channels != 3:
            encoder = encoder.change_input_channels(num_channels)
        
        # Extract 'model_name' and 'pretrained'
        model_name = encoder_config.get('model_name', 'tf_efficientnet_b5.ns_jft_in1k')  # Adjust default if needed
        pretrained = encoder_config.get('pretrained', False)
        
        # Initialize the encoder module
        encoder_module = CustomTimmEncoder(
            model_name=model_name,
            num_channels=num_channels,
            features_only=True,
            pretrained=pretrained
        )
        
        # Adjust encoder input channels if necessary
        encoder_module = adjust_encoder_input_channels(encoder_module, num_channels)
        
        # Extract other configuration parameters
        input_upsample_factor = config.get("input_upsample_factor", 1)
        activation = config.get('activation', 'relu')  # Default to 'relu' if not specified
        num_extra_blocks = config.get('num_extra_blocks', 0)
        num_classes = config.get("num_classes", 2)
        anchors = config.get("anchors", [(10, 13), (16, 30), (33, 23)])
        
        # Extract decoder configuration
        decoder_config = config.decoder
        decoder_features = decoder_config.channels
        block_type = decoder_config.get("block_type", "UnetBlock")
        upsample_type = decoder_config.get("upsample_type", "bilinear")
        
        # Extract neck_out_channels
        neck_out_channels = config.get('neck_out_channels', 256)
        
        # Determine the total scale factor from the decoder
        decoder_scale_factor = 2  # Modify if different in your architecture
        
        # Compute output_stride based on encoder's stride and decoder's upsampling
        output_stride = encoder_module.strides[0] // (decoder_scale_factor * input_upsample_factor)
        
        # Instantiate the box coder
        box_coder = instantiate(
            config.box_coder,
            output_stride=output_stride,
        )
        
        # Define the evaluation function
        def evaluate_architecture(model, train_loader, val_loader):
            return model.evaluate_architecture(model, train_loader, val_loader)
        
        # Instantiate ProgressiveNAS
        nas_config = config.get('nas', {})
        nas = ProgressiveNAS(
            base_model=None,  # Will assign after base_model instantiation
            search_space=nas_config.get('search_space', {}),
            evaluation_function=evaluate_architecture
        )
        
        # Instantiate the base model
        base_model = cls(
            encoder=encoder_module,
            decoder_features=decoder_features,
            heads=heads,
            box_coder=box_coder,
            neck_out_channels=neck_out_channels,
            input_upsample_factor=input_upsample_factor,
            activation=activation,
            block_type=block_type,
            upsample_type=upsample_type,
            num_extra_blocks=num_extra_blocks,
            num_classes=num_classes,
            anchors=anchors,
        )
        
        # Assign the base model to ProgressiveNAS
        nas.base_model = base_model
        
        # Return the instantiated model with NAS
        return cls(
            encoder=encoder_module,
            decoder_features=decoder_features,
            heads=heads,
            box_coder=box_coder,
            nas=nas, 
            neck_out_channels=neck_out_channels,
            input_upsample_factor=input_upsample_factor,
            activation=activation,
            block_type=block_type,
            upsample_type=upsample_type,
            num_extra_blocks=num_extra_blocks,
            num_classes=num_classes,  # Corrected typo here
            anchors=anchors,
        )

def adjust_encoder_input_channels(encoder_module, in_chans):
    # [Existing code...]
    if hasattr(encoder_module, 'encoder'):
        base_encoder = encoder_module.encoder
    else:
        base_encoder = encoder_module

    if hasattr(base_encoder, 'conv_stem'):
        old_conv = base_encoder.conv_stem
     
        new_conv = nn.Conv2d(
            in_channels=in_chans,
            out_channels=old_conv.out_channels,
            kernel_size=old_conv.kernel_size,
            stride=old_conv.stride,
            padding=old_conv.padding,
            bias=old_conv.bias is not None
        )
        
        with torch.no_grad():
            num_existing_in_channels = min(old_conv.in_channels, in_chans)
            new_conv.weight[:, :num_existing_in_channels, ...] = old_conv.weight[:, :num_existing_in_channels, ...]
            if in_chans > old_conv.in_channels:
                torch.nn.init.kaiming_normal_(new_conv.weight[:, old_conv.in_channels:, ...])

        base_encoder.conv_stem = new_conv
        
        print(f"[DEBUG] After adjustment, encoder.conv_stem.in_channels: {base_encoder.conv_stem.in_channels}")
    else:
        raise NotImplementedError("Adjusting input channels for this encoder is not implemented.")
    return encoder_module

class EnhancedCenterNetUNet(nn.Module):
    def __init__(self, channels, num_classes):
        super().__init__()
        self.backbone = CenterNetUNetModel(channels, num_classes)
        self.fpn = FeaturePyramidNetwork([64, 128, 256, 512], 256)
        self.se_blocks = nn.ModuleList([SqueezeExcitation(c) for c in [64, 128, 256, 512]])

    def forward(self, x):
        features = self.backbone(x)
        for i, feature in enumerate(features):
            features[i] = self.se_blocks[i](feature)
        return self.fpn(features)