to be able to give the correct solution you need to see these files 
1-config
2-centernet_unet.py
3-yolo_head.py
4-yolo_decoder.py
5-decoupled_head.py
6-unet_block.py
7-custom_encoder.py
8-necks.py
9-transformer_classifier.py
and before all you should understand that 
This project is a complex integrated multi-model system designed for vessel detection and fishing activity classification using SAR (Synthetic Aperture Radar) TIFF images as input. The system leverages several advanced deep learning architectures and techniques to achieve high accuracy and efficiency in maritime surveillance and monitoring.
Key Components
CenterNet UNet
Combines CenterNet for object detection and U-Net for image segmentation.
YOLO (You Only Look Once)
A real-time object detection system known for its speed and accuracy.
Neural Architecture Search (NAS)
Optimizes model architecture to enhance performance.
Feature Pyramid Network (FPN)
Enables multi-scale feature extraction, improving detection across various object sizes.
Feature Aggregation Network (FAN)
Combines features from different network levels for enriched representation.
Swin Transformer
A hierarchical vision transformer that captures global context effectively.
Optuna
A hyperparameter optimization framework to fine-tune model settings.
Output Format
The system processes SAR images and outputs predictions in a CSV format with the following columns:
vessel_length_m: Length of the detected vessel in meters.
detect_scene_row: Row coordinate in the detection scene.
detect_scene_column: Column coordinate in the detection scene.
scene_id: Identifier for the scene.
is_vessel: Boolean indicating if a vessel is detected.
is_fishing: Boolean indicating if fishing activity is detected.
Sample Output:
vessel_length_m,detect_scene_row,detect_scene_column,scene_id,is_vessel,is_fishing
33.99261,11448,13456,1,True,True
28.166428,15632,12622,1,True,True
32.700775,19454,4840,1,True,True
37.80931,12862,2608,1,True,True
48.392456,4902,5522,1,True,False
The system architecture encompasses model configuration, data augmentation, training processes, and inference pipelines
so please understand all this and remember and respond with done so I can show you the files to deeply review and remember  and respond with only 'done' so we can process with find solution

here's file centernet_unet.py
from functools import partial
from typing import List, Dict, Any, Tuple, Callable
from omegaconf import DictConfig, OmegaConf
from hydra.utils import instantiate, get_class
from pytorch_toolbelt.modules import (
    ACT_RELU,
    EncoderModule,
)
import torchmetrics
from torch.utils.data import DataLoader
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, UpsampleBlock, ResidualDeconvolutionUpsample2d
from .unet_blocks import IRBlock, get_unet_block, SqueezeExcitation, UnetBlock
from .necks import FPNNeck
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.nas import OptunaNAS, FocalLoss
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": YOLOHead,  # Assuming AnchorFreeHead is similar
    # 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 (CustomTimmEncoder): The encoder module.
        in_chans (int): Desired number of input channels.

    Returns:
        CustomTimmEncoder: 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.encoder, name):
            first_conv = getattr(encoder_module.encoder, name)
            break
    
    if first_conv is None:
        if isinstance(encoder_module.encoder, timm.models.SwinTransformer):
            # Transformer-based encoder typically has 'patch_embed.proj'
            if hasattr(encoder_module.encoder.patch_embed, 'proj'):
                first_conv = encoder_module.encoder.patch_embed.proj
            else:
                raise AttributeError(f"The encoder '{type(encoder_module).__name__}' does not have 'patch_embed.proj' attribute.")
        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 hasattr(encoder_module.encoder, 'conv_stem'):
        encoder_module.encoder.conv_stem = new_conv
    elif hasattr(encoder_module.encoder, 'patch_embed'):
        encoder_module.encoder.patch_embed.proj = new_conv
    else:
        setattr(encoder_module.encoder, name, new_conv)
    
    logger.info(f"Adjusted encoder's first conv layer to have {in_chans} input channels.")
    return encoder_module


class CenterNetUNetModel(nn.Module):
    def __init__(
        self,
        config: DictConfig,
        encoder: CustomTimmEncoder,
        decoder_features: List[int],
        heads: List[nn.Module],
        box_coder,
        neck_out_channels: int = 256,
        input_upsample_factor: int = 1,
        activation: str = 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], [30, 61], [62, 45], [59, 119]],
        nas: OptunaNAS = None,
        dropout_rate: float = 0.1
    ):
        super(CenterNetUNetModel, self).__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)]
        )
        downsample_factor = 16
        transformer_img_size = (config.model.image_size[0] // downsample_factor, config.model.image_size[1] // downsample_factor)
        
        # Transformer Encoder Integration
        self.transformer_encoder = SwinTransformerEncoder(
            img_size=transformer_img_size,
            embed_dim=config.model.transformer.embed_dim,
            depths=list(config.model.transformer.depths), 
            num_heads=list(config.model.transformer.num_heads), 
            window_size=tuple(config.model.transformer.window_size), 
            mlp_ratio=config.model.transformer.mlp_ratio,
            qkv_bias=config.model.transformer.qkv_bias,
            qk_scale=config.model.transformer.qk_scale,
            drop_rate=config.model.transformer.drop_rate,
            attn_drop_rate=config.model.transformer.attn_drop_rate,
            drop_path_rate=config.model.transformer.drop_path_rate,
            norm_layer=config.model.transformer.norm_layer,
            in_chans=encoder.channels[-1],
            use_dropout=config.model.transformer.use_dropout,
            use_custom_pos_embed=config.model.transformer.use_custom_pos_embed,
        )
        print(f"[DEBUG] SwinTransformerEncoder initialized with img_size={config.model.image_size} and window_size={config.model.transformer.window_size}")

        # Feature Aggregation Network (FAN)
        self.fan = FeatureAggregationNetwork(
            input_channels=config.model.fan.input_channels,  # [272]
            out_channels=neck_out_channels,                   # 256
            num_layers=config.model.nas.get('num_layers', 4),  # 2
            upsample_mode=config.model.fan.upsample_mode,     # 'bilinear'
            use_attention=config.model.fan.use_attention,     # True
            fusion_method=config.model.fan.fusion_method      # 'weighted_sum'
        )

        # No Channel Alignment Required as FAN takes [272] and outputs [256]

        encoder_channels = encoder.channels[:-1] + [neck_out_channels]  # [24, 48, 64, 160, 256] # [24, 48, 64, 160, 256]

        # Decoder Integration
        unet_block = get_unet_block(block_type, activation=config.model.decoder.activation)
        self.decoder = DebugUnetDecoder(
            encoder_channels=encoder_channels,         # [24, 48, 64, 160, 256]
            decoder_channels=decoder_features,         # [512, 384, 256, 128, 64]
            unet_block=unet_block,
            upsample_block=CenterNetUNetModel.get_upsample_block_static(upsample_type),  # UpsampleBlock with in_channels and out_channels
            center=False,
        )
    
        # Neck Integration
        self.neck = FPNNeck(channels=decoder_features, out_channels=neck_out_channels)  # [512, 384, 256, 128, 64] -> [256]
        self.neck_conv = nn.Conv2d(
            in_channels=neck_out_channels,  # 256
            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=config.model.image_size,  # Adjust based on your image size
            stride=config.model.output_stride,          # 16
            conf_threshold=0.5,
            iou_threshold=0.4
        )

        # Fusion Layer Input Channels
        yolo_out_channels = self.yolo_decoder.out_channels  # Assuming this matches configuration
        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)

        print(f"[DEBUG] CenterNetUNetModel initialized with window_size={config.model.transformer.window_size} and img_size={config.model.image_size}")

    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.UpsamplingBilinear2d,
                scale_factor=2
            )
        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):
        logger.debug("Starting forward pass through CenterNetUNetModel.")
        
        # Initial Upsampling
        logger.debug("Performing initial upsampling.")
        x = self.initial_upsample(x)  # [B, 2, H, W]

        # Encoder Forward Pass
        logger.debug("Passing through encoder.")
        encoder_features = self.encoder(x)

        # Transformer Feature
        transformer_feature = encoder_features[-1]

        # FAN Processing
        aggregated_feature = self.fan([transformer_feature])
        encoder_features[-1] = aggregated_feature

        # Decoder
        logger.debug("Passing aggregated features to decoder.")
        decoded_features = self.decoder(encoder_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
        logger.debug("Passing neck features to heads.")
        head_outputs = {}
        for head in self.heads:
            head_output = head(neck_features)
            if isinstance(head_output, dict):
                head_outputs.update(head_output)
            else:
                raise ValueError("Head output is not a dict.")
        
        # Decode YOLO Outputs (if needed)
        if CENTERNET_OUTPUT_YOLO in head_outputs:
            decoded_yolo = self.yolo_decoder.decode(head_outputs[CENTERNET_OUTPUT_YOLO])
            head_outputs['decoded_yolo'] = decoded_yolo

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

        # Verify Outputs Before Concatenation
        output_keys = [
            CENTERNET_OUTPUT_OBJECTNESS_MAP,
            CENTERNET_OUTPUT_VESSEL_MAP,
            CENTERNET_OUTPUT_FISHING_MAP,
            CENTERNET_OUTPUT_SIZE,
            CENTERNET_OUTPUT_OFFSET,
            CENTERNET_OUTPUT_YOLO,
        ]

        for key in output_keys:
            tensor = outputs.get(key)
            if tensor is None:
                logger.error(f"Output '{key}' is None before concatenation")
                raise ValueError(f"Output '{key}' is None before concatenation")
            else:
                logger.debug(f"Output '{key}' shape: {tensor.shape}")

        # Verify Outputs Before Concatenation
        for idx, key in enumerate([
            CENTERNET_OUTPUT_OBJECTNESS_MAP,
            CENTERNET_OUTPUT_VESSEL_MAP,
            CENTERNET_OUTPUT_FISHING_MAP,
            CENTERNET_OUTPUT_SIZE,
            CENTERNET_OUTPUT_OFFSET,
            CENTERNET_OUTPUT_YOLO,
        ]):
            tensor = outputs[key]
            if tensor is None:
                logger.error(f"Output '{key}' is None before concatenation.")
                raise ValueError(f"Output '{key}' is None before concatenation.")
            else:
                logger.debug(f"Output '{key}' shape: {tensor.shape}")
                
        # Concatenate for Fusion
        tensors_to_concat = [
            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]
        ]
        if CENTERNET_OUTPUT_OFFSET in outputs and outputs[CENTERNET_OUTPUT_OFFSET] is not None:
            tensors_to_concat.append(outputs[CENTERNET_OUTPUT_OFFSET])  # [B, 2, H, W]
        else:
            logger.warning(f"Offset output is missing, skipping in concatenation.")
        tensors_to_concat.append(outputs[CENTERNET_OUTPUT_YOLO])        # [B, YOLO_out_channels, H, W]

        concatenated = torch.cat(tensors_to_concat, dim=1)

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

        # Add fused_output to outputs
        outputs['fused_output'] = fused_output

        return outputs

    def evaluate_architecture(self, model: nn.Module, train_loader: DataLoader, val_loader: DataLoader) -> float:
        """
        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 = self.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_acc, _ = self.evaluate(model, val_loader, torchmetrics.Accuracy().to(device), torchmetrics.F1().to(device))
        print(f"Architecture Evaluation - Validation Accuracy: {val_acc:.4f}")
        return val_acc


    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.

        Calculates the composite loss for multi-task learning with weighting.
        """
        total_loss = 0.0
        losses = {}

        # Define loss weights (can be hyperparameters)
        weights = {
            'heatmap_loss': 1.0,
            'vessel_loss': 1.0,
            'fishing_loss': 1.0,
            'size_loss': 1.0,
            'offset_loss': 1.0,
            'yolo_loss': 1.0,
        }

        # Heatmap loss (Focal Loss)
        if CENTERNET_OUTPUT_OBJECTNESS_MAP in outputs and 'CENTERNET_TARGET_OBJECTNESS_MAP' in targets:
            heatmap_loss = self.heatmap_loss_fn(
                outputs[CENTERNET_OUTPUT_OBJECTNESS_MAP],
                targets['CENTERNET_TARGET_OBJECTNESS_MAP']
            )
            total_loss += weights['heatmap_loss'] * heatmap_loss
            losses['heatmap_loss'] = heatmap_loss.item()

        # Vessel loss (Focal Loss or BCE)
        if CENTERNET_OUTPUT_VESSEL_MAP in outputs and 'CENTERNET_TARGET_VESSEL_MAP' in targets:
            vessel_loss = self.vessel_loss_fn(
                outputs[CENTERNET_OUTPUT_VESSEL_MAP],
                targets['CENTERNET_TARGET_VESSEL_MAP']
            )
            total_loss += weights['vessel_loss'] * vessel_loss
            losses['vessel_loss'] = vessel_loss.item()

        # Fishing loss (Focal Loss or BCE)
        if CENTERNET_OUTPUT_FISHING_MAP in outputs and 'CENTERNET_TARGET_FISHING_MAP' in targets:
            fishing_loss = self.fishing_loss_fn(
                outputs[CENTERNET_OUTPUT_FISHING_MAP],
                targets['CENTERNET_TARGET_FISHING_MAP']
            )
            total_loss += weights['fishing_loss'] * fishing_loss
            losses['fishing_loss'] = fishing_loss.item()

        # Size loss (Smooth L1 Loss)
        if CENTERNET_OUTPUT_SIZE in outputs and 'CENTERNET_TARGET_SIZE' in targets:
            size_loss = F.smooth_l1_loss(
                outputs[CENTERNET_OUTPUT_SIZE],
                targets['CENTERNET_TARGET_SIZE']
            )
            total_loss += weights['size_loss'] * size_loss
            losses['size_loss'] = size_loss.item()

        # Offset loss (Smooth L1 Loss)
        if CENTERNET_OUTPUT_OFFSET in outputs and 'CENTERNET_TARGET_OFFSET' in targets:
            offset_loss = F.smooth_l1_loss(
                outputs[CENTERNET_OUTPUT_OFFSET],
                targets['CENTERNET_TARGET_OFFSET']
            )
            total_loss += weights['offset_loss'] * offset_loss
            losses['offset_loss'] = offset_loss.item()

        # YOLO loss
        if CENTERNET_OUTPUT_YOLO in outputs and CENTERNET_TARGET_YOLO in targets:
            yolo_loss = self.yolo_loss_fn(
                outputs[CENTERNET_OUTPUT_YOLO],
                targets[CENTERNET_TARGET_YOLO]
            )
            total_loss += weights['yolo_loss'] * yolo_loss
            losses['yolo_loss'] = yolo_loss.item()

        # Log individual losses
        for name, value in losses.items():
            logger.debug(f"{name}: {value:.4f}")

        return total_loss


    def evaluate(self, model: nn.Module, val_loader: DataLoader, accuracy_metric: torchmetrics.Metric, f1_metric: torchmetrics.Metric) -> Tuple[float, float]:
        """
        Evaluates the model on the validation set and computes the desired metrics.

        Args:
            model (nn.Module): The model to evaluate.
            val_loader (DataLoader): Validation data loader.
            accuracy_metric (torchmetrics.Metric): Accuracy metric instance.
            f1_metric (torchmetrics.Metric): F1 score metric instance.

        Returns:
            Tuple[float, float]: Validation accuracy and F1 score.
        """
        model.eval()
        accuracy_metric.reset()
        f1_metric.reset()
        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)

                # Compute metrics for vessel classification
                if CENTERNET_OUTPUT_VESSEL_MAP in outputs and 'CENTERNET_TARGET_VESSEL_MAP' in targets:
                    preds = torch.sigmoid(outputs[CENTERNET_OUTPUT_VESSEL_MAP]) > 0.5
                    targets_binary = targets['CENTERNET_TARGET_VESSEL_MAP'] > 0.5
                    accuracy_metric.update(preds, targets_binary)
                    f1_metric.update(preds, targets_binary)

        val_acc = accuracy_metric.compute().item()
        val_f1 = f1_metric.compute().item()
        return val_acc, val_f1

    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():
            if key in [
                'CENTERNET_TARGET_VESSEL_MAP',
                'CENTERNET_TARGET_FISHING_MAP',
                'CENTERNET_TARGET_OBJECTNESS_MAP'
            ]:
                prepared[key] = value.to(dtype=torch.float32)
            elif key in ['CENTERNET_TARGET_SIZE', 'CENTERNET_TARGET_OFFSET']:
                prepared[key] = value.to(dtype=torch.float32)
            elif key == 'CENTERNET_TARGET_YOLO':
                # Ensure YOLO targets are properly formatted tensors
                prepared[key] = value.to(dtype=torch.float32)
            else:
                prepared[key] = value
        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.
        """
        # 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.model.decoder
        decoder_features = decoder_config.decoder_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.model.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.model.encoder
        encoder: CustomTimmEncoder = instantiate(encoder_config)

        # Adjust input channels if necessary
        num_channels = config.get('num_channels', 3)
        if num_channels != 3:
            encoder = adjust_encoder_input_channels(encoder, 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], [30, 61], [62, 45], [59, 119]])

        # Extract decoder configuration
        decoder_config = config.model.decoder
        decoder_features = decoder_config.decoder_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)

        # Extract transformer configuration
        transformer_config = config.model.transformer

        # Instantiate Transformer Encoder
        transformer_encoder = instantiate(transformer_config)

        # Instantiate Feature Aggregation Network (FAN)
        expected_fan_input_channels = config.get('expected_fan_input_channels', 256) 

        # Instantiate Feature Aggregation Network (FAN)
        fan = instantiate(
            config.model.fan,  # Ensure 'fan' is defined in config
            input_channels=[encoder.channels[-1]],  # Corrected to use encoder's output channels
            out_channels=config.model.neck_out_channels
        )

        # Instantiate Decoder
        unet_block = get_unet_block(block_type, activation=config.model.decoder.activation)
        decoder = DebugUnetDecoder(
            encoder_channels=encoder.channels,
            decoder_channels=decoder_features,
            #input_channels=encoder.channels[-1],  # Corrected to use encoder's output channels
            unet_block=unet_block,
            upsample_block=cls.get_upsample_block_static(upsample_type),
            center=False,
        )

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

        # Instantiate Box Coder
        box_coder = MultilabelCircleNetCoder(
            image_size=config.get('image_size', (224,224)),
            output_stride=config.model.get('output_stride', 16),
            max_objects=config.model.box_coder.get('max_objects', 512),
            heatmap_encoding=config.model.box_coder.get('heatmap_encoding', 'umich'),
            labels_encoding=config.model.box_coder.get('labels_encoding', 'circle'),
            ignore_value=config.model.box_coder.get('ignore_value', -1)
        )

        # Instantiate YOLO Decoder
        yolo_decoder = YOLODecoder(
            anchors=anchors,
            num_classes=num_classes,
            input_image_size=config.model.image_size,
            stride=config.model.output_stride,
            conf_threshold=0.5,
            iou_threshold=0.4
        )

        # Instantiate Heads
        instantiated_heads = []
        for head in heads:
            instantiated_heads.append(head)

        # Instantiate NAS if present
        nas = None
        if 'nas' in config:
            nas_config = config.model.nas
            nas = OptunaNAS(
                search_space=nas_config.get('search_space', {}),
                evaluation_function=None,  # Will be set by NAS runner
                max_epochs=nas_config.get('max_epochs', 10),
                use_metrics=nas_config.get('use_metrics', True)
            )

        # Build the model
        model = cls(
            config=config,
            encoder=encoder_module,
            decoder_features=decoder_features,
            heads=instantiated_heads,
            box_coder=box_coder,
            neck_out_channels=neck_out_channels,
            input_upsample_factor=input_upsample_factor,
            activation=activation,
            num_extra_blocks=num_extra_blocks,
            block_type=block_type,
            upsample_type=upsample_type,
            num_classes=num_classes,
            anchors=anchors,
            nas=nas,
            dropout_rate=config.get('dropout_rate', 0.1)
        )

        return model

    
    @staticmethod
    def get_upsample_block_static(upsample_type: str) -> Callable:
        """
        Static 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: Class for the upsample block.
        """
        if upsample_type in ['nearest', 'bilinear']:
            return UpsampleBlock  # Return the custom UpsampleBlock
        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}")
    
    @classmethod
    def instantiate_head(cls, single_head_cfg: Dict[str, Any], decoder_features: List[int]) -> nn.Module:
        """
        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),
                num_classes=single_head_cfg['num_classes'],  # Pass num_classes
                anchors=single_head_cfg['anchors'],          # Pass anchors
            )
        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', (2048, 2048))),
                stride=single_head_cfg.get('stride', 16)
            )
        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):
    """
    Enhanced CenterNetUNet integrating additional components like FPN and SE blocks for improved performance.
    """
    def __init__(self, config: DictConfig, **kwargs):
        super().__init__(config, **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]])

        self.transformer_encoder = SwinTransformerEncoder(
            window_size=tuple(config.model.transformer.window_size),  # Convert list to tuple
            img_size=tuple(config.model.image_size) // config.model.output_stride,
            embed_dim=96,
            depths=[2, 2, 6, 2],
            num_heads=[3, 6, 12, 24],
            mlp_ratio=4.0,
            qkv_bias=True,
            qk_scale=None,
            drop_rate=0.0,
            attn_drop_rate=0.0,
            drop_path_rate=0.1,
            norm_layer=nn.LayerNorm,
            in_chans=encoder.channels[-1],
            use_dropout=False,
            use_custom_pos_embed=False,
        )

    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].unsqueeze(0))
            outputs['enhanced_vessel_map'] = enhanced_vessel.squeeze(0)
            logger.debug("Applied additional FPN to vessel_map.")

        return outputs