DEBUG:rasterio._io:mask_flags: ([<MaskFlags.nodata: 8>],)
DEBUG:rasterio._io:Jump straight to _read()
DEBUG:rasterio._io:Window: Window(col_off=23797, row_off=3633, width=512, height=512)
DEBUG:rasterio._io:IO window xoff=23797.0 yoff=3633.0 width=512.0 height=512.0
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x358435910>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x358435c10> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x358435c10>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x358435910>
[DEBUG] Channel 0 shape: (512, 512)
[DEBUG] Channel 1 shape: (512, 512)
(274.0, 274.0, 1.5899999141693115) [<class 'float'>, <class 'float'>, <class 'float'>]
[DEBUG] Feature 0 shape: torch.Size([16, 24, 256, 256])
[DEBUG] Feature 1 shape: torch.Size([16, 48, 128, 128])
[DEBUG] Feature 2 shape: torch.Size([16, 64, 64, 64])
[DEBUG] Feature 3 shape: torch.Size([16, 160, 32, 32])
[DEBUG] Feature 4 shape: torch.Size([16, 272, 16, 16])
Error executing job with overrides: []
Traceback (most recent call last):
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel3.py", line 806, in <module>
    main()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel3.py", line 802, in main
    pipeline.train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel3.py", line 621, in train
    runner.train(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/runners/runner.py", line 170, in train
    distributed_cmd_run(self.run_experiment, distributed)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/utils/scripts.py", line 176, in distributed_cmd_run
    worker_fn(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/core/runner.py", line 929, in run_experiment
    self._run_event("on_exception")
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/core/runner.py", line 724, in _run_event
    getattr(callback, event)(self)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/callbacks/exception.py", line 34, in on_exception
    raise exception
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/core/runner.py", line 917, in run_experiment
    self._run_stage(stage)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/core/runner.py", line 885, in _run_stage
    self._run_epoch(stage=stage, epoch=self.epoch)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/core/runner.py", line 864, in _run_epoch
    self._run_loader(loader)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/core/runner.py", line 801, in _run_loader
    self._run_batch(batch)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/core/runner.py", line 775, in _run_batch
    self._handle_batch(batch=batch)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/runners/supervised.py", line 144, in _handle_batch
    self.output = self.forward(batch)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/runners/supervised.py", line 131, in forward
    output = self._process_input(batch, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/runners/supervised.py", line 94, in _process_input_str
    output = self.model(batch[self.input_key], **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/centernet_unet.py", line 282, in forward
    transformer_features = self.transformer_encoder(feature_maps[-1])
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/transformer_classifier.py", line 119, in forward
    features = self.swin(x)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/timm/models/swin_transformer.py", line 838, in forward
    x = self.forward_features(x)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/timm/models/swin_transformer.py", line 829, in forward_features
    x = self.patch_embed(x)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/timm/layers/patch_embed.py", line 116, in forward
    _assert(H == self.img_size[0], f"Input height ({H}) doesn't match model ({self.img_size[0]}).")
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/__init__.py", line 1777, in _assert
    assert condition, message
AssertionError: Input height (16) doesn't match model (1024).
1/100 * Epoch (train):   0% 0/27 [00:27<?, ?it/s]
(sam3) samerkharboush@Samers-MBP SAR-MPS-OptunaNasFbnYoloTransformers % 

here's some files might be related

##transformer_classifier.py

from torch import nn
import timm
import torch
from typing import Optional
from timm.models.swin_transformer import SwinTransformer
from xview3.centernet.models.dropath import DropPath
import torchmetrics


class SwinTransformerEncoder(nn.Module):
    def __init__(
        self,
        img_size: int = 1024,
        embed_dim: int = 96,
        depths: list = [2, 2, 6, 2],
        num_heads: list = [3, 6, 12, 24],
        window_size: int = 7,
        mlp_ratio: float = 4.0,
        qkv_bias: bool = True,
        qk_scale: Optional[float] = None,
        drop_rate: float = 0.0,
        attn_drop_rate: float = 0.0,
        drop_path_rate: float = 0.1,
        norm_layer: Optional[nn.Module] = nn.LayerNorm,
        in_chans: int = 2,  # Adjusted to match dataset input channels
        use_dropout: bool = False,
        use_custom_pos_embed: bool = False,
    ):
        """
        Swin Transformer Encoder with enhanced features.

        Args:
            img_size (int): Input image size.
            embed_dim (int): Feature dimension.
            depths (list): Depth of each Swin Transformer stage.
            num_heads (list): Number of attention heads in each stage.
            window_size (int): Window size for attention.
            mlp_ratio (float): Ratio of MLP hidden dim to embedding dim.
            qkv_bias (bool): If True, add a learnable bias to query, key, value.
            qk_scale (Optional[float]): Override default qk scale of head_dim ** -0.5.
            drop_rate (float): Dropout rate.
            attn_drop_rate (float): Attention dropout rate.
            drop_path_rate (float): Stochastic depth rate.
            norm_layer (Optional[nn.Module]): Normalization layer.
            in_chans (int): Number of input channels.
            use_dropout (bool): Whether to include Dropout layers.
            use_custom_pos_embed (bool): Whether to use custom positional embeddings.
        """
        super(SwinTransformerEncoder, self).__init__()
        self.swin = SwinTransformer(
            img_size=img_size,
            patch_size=4,
            in_chans=in_chans,
            embed_dim=embed_dim,
            depths=depths,
            num_heads=num_heads,
            window_size=window_size,
            mlp_ratio=mlp_ratio,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            drop_rate=drop_rate,
            attn_drop_rate=attn_drop_rate,
            drop_path_rate=drop_path_rate,
            norm_layer=norm_layer,
            ape=False,
            patch_norm=True,
            use_checkpoint=False
        )
        
        # Initialize advanced weight initialization
        self._init_weights()
        
        # Integrate TorchMetrics for monitoring
        self.accuracy = torchmetrics.Accuracy(task="multiclass", num_classes=2)
        self.topk_accuracy = torchmetrics.Accuracy(task="multiclass", num_classes=2)
        
        # Optional dropout layer after the transformer
        if use_dropout:
            self.dropout = nn.Dropout(drop_rate)
        else:
            self.dropout = nn.Identity()
        
        self.use_custom_pos_embed = use_custom_pos_embed
        if use_custom_pos_embed:
            self.pos_embed = nn.Parameter(torch.zeros(1, embed_dim, img_size // 4, img_size // 4))
            nn.init.kaiming_normal_(self.pos_embed, mode='fan_out', nonlinearity='relu')

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.trunc_normal_(m.weight, std=.02)
                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)
            elif isinstance(m, nn.Conv2d):
                # check for tensor dimensions
                if m.weight.dim() > 1:
                    nn.init.xavier_uniform_(m.weight)
                else:
                    nn.init.normal_(m.weight)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)


    def forward(self, x):
        """
        Forward pass through the Swin Transformer Encoder.

        Args:
            x (Tensor): Input tensor of shape [B, C, H, W].

        Returns:
            Tensor: Extracted features from the encoder.
        """
        if self.use_custom_pos_embed:
            x = x + self.pos_embed
        features = self.swin(x)
        features = self.dropout(features)
        return features


###centernet_unet.py

from functools import partial
from typing import List, Dict, Any, Tuple
from omegaconf import DictConfig , OmegaConf
from hydra.utils import instantiate, get_class
from pytorch_toolbelt.modules import (
    ACT_RELU,
    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 OptunaNAS, FocalLoss
from .heads.yolo_head import YOLOHead, AnchorFreeHead
from .heads.yolo_decoder import YOLODecoder
from xview3.centernet.models.transformer_classifier import SwinTransformerEncoder
from xview3.centernet.models.necks import FPNNeck , FeatureAggregationNetwork
from xview3.centernet.models.heads import DecoupledHeadGroupNormLateShuffle
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
import logging
from xview3.utils import TorchDevice
from omegaconf.errors import ConfigKeyError  # Added import for ConfigKeyError
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)


__all__ = ["CenterNetUNetModel"]

device = TorchDevice.choose_torch_device()


HEAD_REGISTRY = {
    "xview3.centernet.models.heads.DecoupledHeadGroupNormLateShuffle": DecoupledHeadGroupNormLateShuffle,
    "xview3.centernet.models.heads.yolo_head.YOLOHead": YOLOHead,
    "xview3.centernet.models.heads.AnchorFreeHead": AnchorFreeHead,
    # Add other head types here
}

def adjust_encoder_input_channels(encoder_module, in_chans):
    """
    Adjusts the input channels of the encoder's first convolutional layer to match the desired number of input channels.
    
    Args:
        encoder_module (nn.Module): The encoder module.
        in_chans (int): Desired number of input channels.

    Returns:
        nn.Module: Encoder module with adjusted input channels.
    """
    first_conv = None
    # Search for common first conv layers
    possible_names = ['conv_stem', 'conv1', 'initial_conv']
    for name in possible_names:
        if hasattr(encoder_module, name):
            first_conv = getattr(encoder_module, name)
            break
    
    if first_conv is None:
        if isinstance(encoder_module, CustomTimmEncoder):
            # Handle CustomTimmEncoder specifically
            first_conv = encoder_module.encoder.conv_stem
        else:
            raise NotImplementedError(f"Adjusting input channels for encoder '{type(encoder_module).__name__}' is not implemented. Expected one of {possible_names}.")
    
    old_conv = first_conv
    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:, ...])

    if isinstance(encoder_module, CustomTimmEncoder):
        encoder_module.encoder.conv_stem = new_conv
    else:
        setattr(encoder_module, name, new_conv)
    
    logger.info(f"Adjusted encoder's first conv layer '{name}' to have {in_chans} input channels.")
    return encoder_module


class CenterNetUNetModel(nn.Module):
    """
    CenterNet-UNet Model integrating an encoder, transformer encoder, feature aggregation network, decoder, neck, and multiple heads.

    Args:
        encoder (EncoderModule): The backbone encoder module.
        decoder_features (List[int]): List of feature channels for the decoder.
        heads (List[nn.Module]): List of head modules (e.g., DecoupledHead, YOLOHead).
        box_coder: The box coder for decoding CenterNet outputs.
        neck_out_channels (int, optional): Output channels for the neck module. Defaults to 256.
        input_upsample_factor (int, optional): Upsampling factor for initial input. Defaults to 1.
        activation (str, optional): Activation function. Defaults to ACT_RELU.
        num_extra_blocks (int, optional): Number of extra encoder blocks. Defaults to 0.
        block_type (str, optional): Type of U-Net block to use. Defaults to "UnetBlock".
        upsample_type (str, optional): Type of upsampling method. Defaults to "bilinear".
        extra_block (Callable, optional): Extra block to add to the encoder. Defaults to partial(IRBlock, stride=2).
        num_classes (int, optional): Number of target classes. Defaults to 2.
        anchors (List[Tuple[int, int]], optional): List of anchor box dimensions for YOLO. Defaults to [(10, 13), (16, 30), (33, 23)].
        nas (OptunaNAS, optional): NAS module for architecture search. Defaults to None.
        dropout_rate (float, optional): Dropout rate. Defaults to 0.1.
        num_layers (int, optional): Number of layers for NAS. Defaults to 4.
        num_filters (int, optional): Number of filters for NAS. Defaults to 256.
        kernel_size (int, optional): Kernel size for NAS. Defaults to 3.
    """
    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: int = 0,
        block_type: str = "UnetBlock",
        upsample_type: str = "bilinear",
        extra_block=partial(IRBlock, stride=2),
        num_classes: int = 2,
        anchors: List[Tuple[int, int]] = [(10, 13), (16, 30), (33, 23)],
        nas: OptunaNAS = None,
        dropout_rate: float = 0.1,
        num_layers: int = 4,        
        num_filters: int = 256,     
        kernel_size: int = 3        
    ):
        super().__init__()

        # Initial Upsampling
        self.initial_upsample = (
            nn.UpsamplingBilinear2d(scale_factor=input_upsample_factor)
            if input_upsample_factor != 1 else nn.Identity()
        )

        # Encoder Forward Pass
        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 Integration
        self.transformer_encoder = SwinTransformerEncoder(
            img_size=encoder.image_size,  # Fetch dynamically
            embed_dim=encoder.embed_dim,  # Fetch dynamically
            depths=encoder.depths,        # Fetch dynamically
            num_heads=encoder.num_heads,  # Fetch dynamically
            window_size=encoder.window_size,  # Fetch dynamically
            mlp_ratio=encoder.mlp_ratio, 
            qkv_bias=encoder.qkv_bias,
            qk_scale=encoder.qk_scale,
            drop_rate=encoder.drop_rate,
            attn_drop_rate=encoder.attn_drop_rate,
            drop_path_rate=encoder.drop_path_rate,
            norm_layer=encoder.norm_layer,
            in_chans=encoder.channels[0],  
            use_dropout=encoder.use_dropout,
            use_custom_pos_embed=encoder.use_custom_pos_embed,
        )

        # Feature Aggregation Network (FAN)
        self.fan = FeatureAggregationNetwork(
            input_channels=[encoder.transformer_output_channels], 
            output_channels=neck_out_channels,
            num_layers=num_layers,
            num_filters=num_filters,
            kernel_size=kernel_size,
            activation=activation,
            dropout_rate=dropout_rate
        )

        # Decoder Integration
        unet_block = get_unet_block(block_type, activation=activation)
        self.decoder = DebugUNetDecoder(
            encoder_channels=encoder.channels,
            decoder_channels=decoder_features,
            input_channels=encoder.transformer_output_channels, 
            unet_block=unet_block,
            upsample_block=self.get_upsample_block(upsample_type),
            center=False,
        )

        # Neck Integration
        self.neck = FPNNeck(channels=decoder_features, out_channels=neck_out_channels)
        self.neck_conv = nn.Conv2d(
            in_channels=neck_out_channels,  # Output from FPNNeck
            out_channels=neck_out_channels,
            kernel_size=1
        )
        self.batch_norm = FilterResponseNormalization(neck_out_channels)
        self.dropout = nn.Dropout(p=0.5)

        # Heads Integration
        self.heads = nn.ModuleList(heads)  # Initialize heads from configuration
        self.head_output_channels = sum([head.out_channels for head in self.heads])  # Sum of all head output channels
        self.output_channels = self.head_output_channels

        self.nas = nas  # Initialize NAS if provided

        # Initialize YOLO Decoder
        self.yolo_decoder = YOLODecoder(
            anchors=anchors,
            num_classes=num_classes,
            input_image_size=encoder.image_size,  # Adjust based on your image size
            stride=encoder.output_stride,         # Adjust based on encoder
            conf_threshold=0.5,
            iou_threshold=0.4
        )

        # Fusion Layer Input Channels
        yolo_out_channels = self.yolo_decoder.out_channels
        fusion_input_channels = 1 + 1 + 1 + 1 + 2 + yolo_out_channels 
        self.fusion = nn.Conv2d(fusion_input_channels, self.head_output_channels, kernel_size=1)
        self.dropout_fusion = nn.Dropout(p=0.5)


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

    def get_upsample_block(self, upsample_type: str):
        """
        Factory method to get the appropriate upsample block based on the upsample_type.

        Args:
            upsample_type (str): Type of upsampling ('nearest', 'bilinear', 'rdtsc', 'shuffle').

        Returns:
            Callable: Partial function for the upsample block.
        """
        if upsample_type in ['nearest', 'bilinear']:
            return partial(
                nn.Upsample,
                scale_factor=2,
                mode=upsample_type,
                align_corners=False
            )
        elif upsample_type == 'rdtsc':
            return partial(ResidualDeconvolutionUpsample2d, scale_factor=2)
        elif upsample_type == 'shuffle':
            return partial(nn.PixelShuffle, upscale_factor=2)
        else:
            raise ValueError(f"Unknown upsample_type {upsample_type}")

    def forward(self, x):
        # Initial Upsampling
        x = self.initial_upsample(x)

        # Encoder Forward Pass
        feature_maps = self.encoder(x)
        for block in self.extra_encoder_blocks:
            feature_maps[-1] = block(feature_maps[-1])

        # Transformer Encoder Forward Pass
        transformer_features = self.transformer_encoder(feature_maps[-1])

        # Feature Aggregation
        aggregated_features = self.fan(transformer_features)

        # Decoder Forward Pass
        decoded_features = self.decoder(aggregated_features)

        # 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 Forward Pass
        head_outputs = {}
        for head in self.heads:
            head_output = head(neck_features)
            if isinstance(head, DecoupledHeadGroupNormLateShuffle):
                head_outputs.update(head_output)
            elif isinstance(head, YOLOHead):
                head_outputs[CENTERNET_OUTPUT_YOLO] = head_output
            else:
                logger.warning(f"Unhandled head type: {type(head)}")

        # Decode YOLO Outputs
        decoded_yolo = self.yolo_decoder.decode(head_outputs.get(CENTERNET_OUTPUT_YOLO))
        head_outputs['decoded_yolo'] = decoded_yolo

        # Extract Standard Outputs
        outputs = {
            CENTERNET_OUTPUT_OBJECTNESS_MAP: head_outputs.get(CENTERNET_OUTPUT_OBJECTNESS_MAP),
            CENTERNET_OUTPUT_SIZE: head_outputs.get(CENTERNET_OUTPUT_SIZE),
            CENTERNET_OUTPUT_OFFSET: head_outputs.get(CENTERNET_OUTPUT_OFFSET),
            CENTERNET_OUTPUT_VESSEL_MAP: head_outputs.get(CENTERNET_OUTPUT_VESSEL_MAP),
            CENTERNET_OUTPUT_FISHING_MAP: head_outputs.get(CENTERNET_OUTPUT_FISHING_MAP),
            CENTERNET_OUTPUT_YOLO: head_outputs.get(CENTERNET_OUTPUT_YOLO),
        }

        # Concatenate All Relevant Feature Maps
        concatenated = torch.cat([
            outputs[CENTERNET_OUTPUT_OBJECTNESS_MAP],  # [B, 1, H, W]
            outputs[CENTERNET_OUTPUT_VESSEL_MAP],      # [B, 1, H, W]
            outputs[CENTERNET_OUTPUT_FISHING_MAP],     # [B, 1, H, W]
            outputs[CENTERNET_OUTPUT_SIZE],            # [B, 1, H, W]
            outputs[CENTERNET_OUTPUT_OFFSET],          # [B, 2, H, W]
            outputs[CENTERNET_OUTPUT_YOLO],            # [B, YOLO_out_channels, H, W]
        ], dim=1)

        # **Verification Prints:**
        logger.debug(f"objectness_map shape: {outputs[CENTERNET_OUTPUT_OBJECTNESS_MAP].shape}")      
        logger.debug(f"vessel_map shape: {outputs[CENTERNET_OUTPUT_VESSEL_MAP].shape}")                  
        logger.debug(f"fishing_map shape: {outputs[CENTERNET_OUTPUT_FISHING_MAP].shape}")                
        logger.debug(f"size_map shape: {outputs[CENTERNET_OUTPUT_SIZE].shape}")                      
        logger.debug(f"offset_map shape: {outputs[CENTERNET_OUTPUT_OFFSET].shape}")                  
        logger.debug(f"yolo_output shape: {head_outputs[CENTERNET_OUTPUT_YOLO].shape}")                
        logger.debug(f"Total concatenated channels: {concatenated.shape[1]}")  

        # **Assertion for Debugging**
        expected_channels = 1 + 1 + 1 + 1 + 2 + head_outputs[CENTERNET_OUTPUT_YOLO].shape[1]
        assert concatenated.shape[1] == expected_channels, \
            f"Expected {expected_channels} channels, but got {concatenated.shape[1]}"

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

        # Optionally, add fused_output to outputs if needed for further processing
        outputs['fused_output'] = fused_output

        return outputs

    def evaluate_architecture(self, model, train_loader, val_loader):
        """
        Defines how to evaluate the architecture. This method trains the model for a few epochs and returns a validation metric.

        Args:
            model (nn.Module): The model to evaluate.
            train_loader (DataLoader): Training data loader.
            val_loader (DataLoader): Validation data loader.

        Returns:
            float: Validation metric (e.g., Mean Average Precision).
        """
        # Define training parameters for architecture evaluation
        optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
        criterion = FocalLoss(alpha=2.0, beta=4.0, eps=1e-6)
        num_epochs = 3  # Train for a few epochs to evaluate architecture

        model.to(device)
        model.train()

        for epoch in range(num_epochs):
            running_loss = 0.0
            for batch in train_loader:
                inputs, targets = self.get_batch(batch)
                inputs = inputs.to(device)
                targets = self.prepare_targets(targets).to(device)

                optimizer.zero_grad()

                outputs = model(inputs)
                loss = model.calculate_loss(outputs, targets)
                loss.backward()
                optimizer.step()

                running_loss += loss.item()

            avg_loss = running_loss / len(train_loader)
            logger.info(f"Architecture Evaluation Epoch {epoch + 1}/{num_epochs} - Loss: {avg_loss:.4f}")

        # Compute validation metric
        val_metric = self.evaluate(model, val_loader)
        logger.info(f"Architecture Evaluation - Validation Metric: {val_metric:.4f}")
        return val_metric

    def calculate_loss(self, outputs: Dict[str, torch.Tensor], targets: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        Calculates the composite loss for multi-task learning.

        Args:
            outputs (Dict[str, torch.Tensor]): Model outputs.
            targets (Dict[str, torch.Tensor]): Ground truth targets.

        Returns:
            torch.Tensor: Total loss.
            
        """
        if CENTERNET_TARGET_YOLO not in targets:
            raise KeyError(f"Expected key '{CENTERNET_TARGET_YOLO}' not found in targets.")

        # CircleNet Loss Components
        circle_net_loss = self.circle_net_loss(**outputs, **targets)

        # YOLO Loss
        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 Calculation
        total_loss = (
            circle_net_loss['objectness'] +
            circle_net_loss['vessel'] +
            circle_net_loss['fishing'] +
            circle_net_loss['size'] +
            yolo_loss
        )
        return total_loss

    def evaluate(self, model: nn.Module, val_loader) -> float:
        """
        Evaluates the model on the validation set and computes the desired metric.

        Args:
            model (nn.Module): The model to evaluate.
            val_loader (DataLoader): Validation data loader.

        Returns:
            float: Validation metric (e.g., Mean Average Precision).
        """
        model.eval()
        metric = 0.0
        with torch.no_grad():
            for batch in val_loader:
                inputs, targets = self.get_batch(batch)
                inputs = inputs.to(device)
                targets = self.prepare_targets(targets).to(device)

                outputs = model(inputs)
                current_metric = self.compute_metric(outputs, targets)
                metric += current_metric

        # Average metric over the validation set
        metric /= len(val_loader)
        return metric

    def compute_metric(self, outputs: Dict[str, torch.Tensor], targets: Dict[str, torch.Tensor]) -> float:
        """
        Computes the desired metric between outputs and targets.

        Args:
            outputs (Dict[str, torch.Tensor]): Model outputs.
            targets (Dict[str, torch.Tensor]): Ground truth targets.

        Returns:
            float: Computed metric (e.g., mAP, IoU).
        """
        # Placeholder: Binary accuracy for vessel classification
        if 'vessel_map' in outputs and 'CENTERNET_TARGET_VESSEL_MAP' in targets:
            preds = torch.sigmoid(outputs['vessel_map']) > 0.5
            target = targets['CENTERNET_TARGET_VESSEL_MAP'] > 0.5
            correct = (preds == target).float().sum()
            total = target.numel()
            accuracy = correct / total
            return accuracy.item()
        else:
            # Default fallback metric
            return 0.0

    def get_batch(self, batch: Tuple[Any, Any]) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """
        Processes a batch from the DataLoader.

        Args:
            batch (Tuple[Any, Any]): A batch from the DataLoader.

        Returns:
            Tuple[torch.Tensor, Dict[str, torch.Tensor]]: Processed inputs and prepared targets.
        """
        inputs, targets = batch
        return inputs, targets

    def prepare_targets(self, targets: Dict[str, Any]) -> Dict[str, torch.Tensor]:
        """
        Prepares targets for the loss function.

        Args:
            targets (Dict[str, Any]): Raw targets from the DataLoader.

        Returns:
            Dict[str, torch.Tensor]: Prepared targets as tensors.
        """
        prepared = {}
        for key, value in targets.items():
            prepared[key] = value.to(dtype=torch.float32)
        return prepared

    @classmethod
    def from_config(cls, config: DictConfig):
        """
        Instantiates the CenterNetUNetModel from a configuration dictionary.

        Args:
            config (DictConfig): Configuration dictionary.

        Returns:
            CenterNetUNetModel: The instantiated model.
        """
        # Replace with appropriate import if needed

        # 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 = []

        # Handling Multiple Heads
        for head_name, single_head_cfg in head_cfg.items():
            head = cls.instantiate_head(single_head_cfg, decoder_features)
            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', 'swin_base_patch4_window7_224')  # Adjust default if needed
        pretrained = encoder_config.get('pretrained', True)

        # Initialize the encoder module
        encoder_module = CustomTimmEncoder(
            model_name=model_name,
            num_channels=num_channels,
            features_only=True,
            pretrained=pretrained,
            image_size=config.get('image_size', (224, 224))
        )

        # 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,
        )

        # Instantiate OptunaNAS if present in config
        nas = None
        if 'nas' in config:
            nas_config = config.nas
            # Instantiate the base model first without NAS
            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,
                nas=None,  # Temporarily set NAS to None
                dropout_rate=config.get('dropout_rate', 0.1),
                num_layers=config.get('num_layers', 4),
                num_filters=config.get('num_filters', 256),
                kernel_size=config.get('kernel_size', 3)
            )

            # Define the evaluation function
            def evaluate_architecture(trial_model, train_loader, val_loader):
                return cls.evaluate_architecture(cls, trial_model, train_loader, val_loader)

            # Instantiate OptunaNAS with the evaluation function
            nas = OptunaNAS(
                search_space=nas_config.get('search_space', {}),
                evaluation_function=lambda trial_model, train_loader, val_loader: cls.evaluate_architecture(trial_model, train_loader, val_loader),
                max_epochs=nas_config.get('max_epochs', 10),
                use_metrics=nas_config.get('use_metrics', True)
            )

            # Assign the base model to OptunaNAS
            nas.base_model = base_model

            # Re-instantiate the base model with NAS
            base_model.nas = nas

            # Return the base model with NAS
            return base_model

        # Instantiate the base model without NAS
        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,
            nas=nas,
            dropout_rate=config.get('dropout_rate', 0.1),
            num_layers=config.get('num_layers', 4),
            num_filters=config.get('num_filters', 256),
            kernel_size=config.get('kernel_size', 3)
        )

        return base_model


    @staticmethod
    def instantiate_head(single_head_cfg: Dict[str, Any], decoder_features: List[int]):
        """
        Instantiates a head module based on the head configuration.

        Args:
            single_head_cfg (Dict[str, Any]): Configuration dictionary for the head.
            decoder_features (List[int]): List of decoder feature channels.

        Returns:
            nn.Module: Instantiated head module.
        """
        head_type = single_head_cfg.get('type', single_head_cfg.get('_target_', None))
        if head_type is None:
            raise ConfigKeyError("Both 'type' and '_target_' are missing from the head configuration")

        HeadClass = HEAD_REGISTRY.get(head_type, None)
        if HeadClass is None:
            raise NotImplementedError(f"Head type '{head_type}' is not registered.")

        if head_type == 'xview3.centernet.models.heads.DecoupledHeadGroupNormLateShuffle':
            head = HeadClass(
                channels=single_head_cfg.get('channels', decoder_features[-1]),
                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'),
                num_blocks=single_head_cfg.get('num_blocks', 3),
            )
        elif head_type == 'xview3.centernet.models.heads.yolo_head.YOLOHead':
            head = HeadClass(
                num_classes=single_head_cfg['num_classes'],
                anchors=single_head_cfg['anchors'],
                input_image_size=tuple(single_head_cfg.get('input_image_size', (1024, 1024))),
                stride=single_head_cfg.get('stride', 32)
            )
        elif head_type == 'xview3.centernet.models.heads.AnchorFreeHead':
            head = HeadClass(
                in_channels=decoder_features[-1],
                num_classes=single_head_cfg['num_classes']
            )
        else:
            raise NotImplementedError(f"Head type '{head_type}' instantiation is not implemented.")

        return head

    def evaluate_architecture(self, model, train_loader, val_loader):
        """
        Defines how to evaluate the architecture. This method trains the model for a few epochs and returns a validation metric.

        Args:
            model (nn.Module): The model to evaluate.
            train_loader (DataLoader): Training data loader.
            val_loader (DataLoader): Validation data loader.

        Returns:
            float: Validation metric (e.g., Mean Average Precision).
        """
        # Define training parameters for architecture evaluation
        optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
        num_epochs = 3  # Train for a few epochs to evaluate architecture

        model.to(device)
        model.train()

        for epoch in range(num_epochs):
            running_loss = 0.0
            for batch in train_loader:
                inputs, targets = self.get_batch(batch)
                inputs = inputs.to(device)
                targets = self.prepare_targets(targets).to(device)

                optimizer.zero_grad()

                outputs = model(inputs)
                loss = model.calculate_loss(outputs, targets)
                loss.backward()
                optimizer.step()

                running_loss += loss.item()

            avg_loss = running_loss / len(train_loader)
            print(f"Architecture Evaluation Epoch {epoch + 1}/{num_epochs} - Loss: {avg_loss:.4f}")

        # Compute validation metric
        val_metric = self.evaluate(model, val_loader)
        print(f"Architecture Evaluation - Validation Metric: {val_metric:.4f}")
        return val_metric

class EnhancedCenterNetUNet(CenterNetUNetModel):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # Initialize additional components
        self.fpn_extra = FeaturePyramidNetwork([64, 128, 256, 512], 256)
        self.se_blocks = nn.ModuleList([SqueezeExcitation(c) for c in [64, 128, 256, 512]])

    def forward(self, x):
        # Call the parent forward
        outputs = super().forward(x)

        # Extract feature maps for additional FPN and SE processing
        feature_maps = list(outputs.values())  # Adjust based on actual feature map routing
        for i, feature in enumerate(feature_maps):
            if i < len(self.se_blocks):
                feature = self.se_blocks[i](feature)
            feature_maps[i] = self.fpn_extra(feature)

        # Update outputs with enhanced feature maps
        for idx, key in enumerate([CENTERNET_OUTPUT_OBJECTNESS_MAP, CENTERNET_OUTPUT_VESSEL_MAP,
                                   CENTERNET_OUTPUT_FISHING_MAP, CENTERNET_OUTPUT_SIZE]):
            if key in outputs and idx < len(self.se_blocks):
                outputs[key] = self.se_blocks[idx](outputs[key])
                logger.debug(f"Applied SE block to {key}.")

        # Optionally, process with additional FPN
        # For example, enhance vessel_map using FPN
        if CENTERNET_OUTPUT_VESSEL_MAP in outputs:
            enhanced_vessel = self.fpn_extra([outputs[CENTERNET_OUTPUT_VESSEL_MAP]])
            outputs['enhanced_vessel_map'] = enhanced_vessel[0]
            logger.debug("Applied additional FPN to vessel_map.")

        return outputs