please can you modify the expected input[8, 1280, 256, 256] to have 1024 channels, 

as I face this error which needs you to understand some files to correctly fix,
  return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Given groups=1, weight of size [256, 1024, 1, 1], expected input[8, 1280, 256, 256] to have 1024 channels, but got 1280 channels instead

first file you need to deeply understand 
#multilabel_circle_net.yaml

defaults:
  - _self_
  - dataset: valid_sar_cv
  - augs: flips_light
  - sampler: balanced
  - normalization: default
  - loss: centernet/rfl_bce_mse
  - optimizer: adamw
  - scheduler: cos
  - train: default
  - runner: default
  - ema: disabled
  - torch: default

batch_size: 8
channels_last: False

local_rank: ${oc.env:LOCAL_RANK, 0}
world_size: ${oc.env:WORLD_SIZE, 0}

experiment: null
seed: 42
deterministic: False
detect_anomaly: False
find_unused: False

transfer: null
checkpoint: null

nas:
  search_space:
    num_layers: [2, 5]
    num_filters: [32, 128]
    kernel_size: [3, 5]
    activation: ['relu', 'leaky_relu']
  num_iterations: 10

model:
  _target_: xview3.centernet.models.centernet_unet.CenterNetUNetModel.from_config
  num_channels: 2
  num_classes: 2
  anchors:
    - [10, 13]
    - [16, 30]
    - [33, 23]
  input_upsample_factor: 1
  encoder:
    _target_: custom_encoder.CustomTimmEncoder
    model_name: 'tf_efficientnet_b5_ns'
    num_channels: 2
    features_only: True
    pretrained: False
  decoder:
    channels: [512, 384, 256, 128]
    block_type: UnetBlock
    upsample_type: bilinear
  head:
    _target_: xview3.centernet.models.heads.decoupled_head.DecoupledHeadGroupNormLateShuffle
    channels: 32
    classifier_dim: 64
    dropout_rate: 0.1
  box_coder:
    _target_: xview3.centernet.bboxer.MultilabelCircleNetCoder
    image_size: [768, 768]
    output_stride: 4
    max_objects: 512
    heatmap_encoding: umich
    labels_encoding: circle
    ignore_value: -1
  activation: silu
  num_extra_blocks: 0

yolo:
  enabled: true
  loss_weight: 1.0

transformer:
  enabled: true
  loss_weight: 1.0

train:
  epochs: 100
  early_stopping: 25
  train_image_size: [768, 768]
  valid_image_size: [2048, 2048]

torch:
  cudnn:
    allow_tf32: True
  cuda_allow_tf32: True
  detect_anomaly: False
  deterministic: False
  benchmark: True

hydra:
  output_subdir: null
  run:
    dir: .
  sweep:
    dir: .
    subdir: .

augs:
  spatial:
    - albumentations.HorizontalFlip
    - albumentations.VerticalFlip
  individual:
    vv0:
      _target_: xview3.augmentations.unclipped_random_brightness.UnclippedRandomBrightnessContrast
      always_apply: false
      p: 0.5

normalization:
  channels:
    vv:
      _target_: xview3.SigmoidNormalization
      midpoint: -20.0
      temperature: 0.18
    vh:
      _target_: xview3.SigmoidNormalization
      midpoint: -20.0
      temperature: 0.18
    bathymetry:
      _target_: xview3.CubicRootNormalization
    wind_direction:
      _target_: albumentations.NoOp
    wind_speed:
      _target_: albumentations.NoOp
    mask:
      _target_: albumentations.NoOp

dataset:
  data_dir: ${oc.env:XVIEW3_DIR}/full_dataset
  splitter: valid_sar_cv
  fold: 0
  num_folds: 4

#next file nas.py
import sys
import traceback
import warnings
from typing import Dict, Any, OrderedDict, Tuple, Optional, List
import catalyst
import os
import numpy as np
import pandas as pd
import torch
import hydra
from torch.optim import Optimizer
from torch.utils.data import Dataset, Sampler, DataLoader
from torch import nn
from omegaconf import DictConfig
from functools import partial
import albumentations as A
from catalyst.callbacks import ControlFlowCallback, CriterionCallback, SchedulerCallback, TimerCallback
from catalyst.core import Callback
from pytorch_toolbelt.utils.catalyst import (
    ShowPolarBatchesCallback,
    BestMetricCheckpointCallback,
    HyperParametersCallback,
    MixupCriterionCallback,
    MixupInputCallback,
)
from pytorch_toolbelt.optimization.functional import freeze_model, get_lr_decay_parameters, get_optimizable_parameters
from pytorch_toolbelt.utils import fs, transfer_weights, count_parameters
from pytorch_toolbelt.utils.distributed import is_main_process
from xview3.centernet.bboxer import MultilabelCircleNetCoder
from xview3.early_stopping import EarlyStoppingCallback
from xview3.nas import ProgressiveNAS
from xview3.centernet.models.heads.yolo_head import YOLOHead
from xview3.centernet.models.transformer_classifier import ConvolutionalClassifier
from xview3.centernet.losses.circlenet_loss import CircleNetLoss, YOLOLoss
from xview3.centernet import *
from xview3.pipeline import *
from xview3.factory import *
from xview3.augmentations import *
from xview3.samplers import *
from xview3.dataset import *
from xview3.ema import *
from xview3.evaluation import *
from xview3.inference import *
from xview3.metric import *
from xview3.constants import INPUT_SCENE_ID_KEY
import logging

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

def warn_with_traceback(message, category, filename, lineno, file=None, line=None):
    log = file if hasattr(file, "write") else sys.stderr
    traceback.print_stack(file=log)
    log.write(warnings.formatwarning(message, category, filename, lineno))

warnings.showwarning = warn_with_traceback

if 'XVIEW3_DIR' not in os.environ:
    os.environ['XVIEW3_DIR'] = '/Volumes/Crucial/xview3_dataset'

class CenternetPipeline(Pipeline):
    def __init__(self, cfg):
        super().__init__(cfg)
        self.box_coder: Optional[MultilabelCircleNetCoder] = None
        self.shore_root = os.path.join(os.environ.get('XVIEW3_DIR', ''), "full_dataset", "shoreline")

    def build_datasets(self, config) -> Tuple[Dataset, Dataset, Optional[Sampler], List[Callback]]:
        dataset = MultilabelCircleNetDataModule(data_dir=os.path.join(os.environ.get('XVIEW3_DIR', '/Volumes/Crucial/xview3_dataset'), 'full_dataset'))
        self.box_coder = MultilabelCircleNetCoder(
            image_size=(config.train.train_image_size[0], config.train.train_image_size[1]),
            output_stride=4,  # Adjust this value as needed
            max_objects=512,  # Adjust this value as needed
            heatmap_encoding="umich",
            labels_encoding="circle",
            ignore_value=-1
        )
        print(f"[DEBUG] self.box_coder in build_datasets: {self.box_coder}")
        if self.box_coder is None:
            raise ValueError("self.box_coder is None in build_datasets.")

        train_df, valid_df, _, shore_root = dataset.train_val_split(
            splitter=config.dataset.splitter,
            fold=config.dataset.fold,
            num_folds=config.dataset.num_folds,
        )

        self.valid_df = valid_df.copy()  # Keep a copy of unmodified groundtruth

        self.master_print("Dataset", config.dataset.slug)
        self.master_print("\tignore_low_confidence_detections", config.dataset.ignore_low_confidence_detections)
        self.master_print("\tignore_low_confidence_labels    ", config.dataset.ignore_low_confidence_labels)
        self.master_print("\tfilter_low_confidence_objects   ", config.dataset.filter_low_confidence_objects)

        self.master_print("Train dataset")
        self.master_print("\tUnknown vessels", np.isnan(train_df.is_vessel.values.astype(float)).sum())
        self.master_print("\tUnknown fishing", np.isnan(train_df.is_fishing.values.astype(float)).sum())
        self.master_print("Valid dataset")
        self.master_print("\tUnknown vessels", np.isnan(valid_df.is_vessel.values.astype(float)).sum())
        self.master_print("\tUnknown fishing", np.isnan(valid_df.is_fishing.values.astype(float)).sum())

        if config.dataset.filter_low_confidence_objects and config.dataset.ignore_low_confidence_labels:
            raise ValueError(f"filter_low_confidence_objects and ignore_low_confidence_labels are mutually exclusive")

        if config.dataset.filter_low_confidence_objects:
            train_df = filter_low_confidence_objects(train_df)
            valid_df = filter_low_confidence_objects(valid_df)
            self.master_print("Excluding low-confidence objects from training & validation")
        elif config.dataset.ignore_low_confidence_labels:
            train_df = ignore_low_confidence_objects(train_df)
            valid_df = ignore_low_confidence_objects(valid_df)
            self.master_print("Ignoring low-confidence labels from training & validation")

        self.shore_root = os.path.join(os.environ.get('XVIEW3_DIR', ''), "full_dataset", "shoreline")

        normalization = build_normalization(config.normalization)

        augmentations = [
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.5),
            A.RandomRotate90(p=0.5),
            A.Resize(height=config.train.train_image_size[0], width=config.train.train_image_size[1], always_apply=True),
        ]

        transform = A.Compose(
            augmentations,
            keypoint_params=A.KeypointParams(
                format='xys',
                label_fields=['confidences', 'is_vessel', 'is_fishing', 'is_near_shore'],
                remove_invisible=False
            )
        )


        individual_augmentations = dict(
            [(key, A.Compose(build_augmentations([value]))) for key, value in config.augs.individual.items() if value is not None]
        )

        for channel in config.dataset.channels:
            if channel not in normalization:
                normalization[channel] = A.NoOp()

        if config.sampler.sampler_type == "each_object":
            train_ds, train_sampler = dataset.get_random_crop_each_object_dataset(
                train_df,
                train_image_size=config.train.train_image_size,
                input_channels=config.dataset.channels,
                box_coder=self.box_coder,  # Pass `self.box_coder` here
                num_samples=config.sampler.num_samples,
                individual_augmentations=individual_augmentations,
                augmentations=augmentations,
                transform=transform,
                normalization=normalization,
                balance_crowd=config.sampler.balance_crowd,
                balance_near_shore=config.sampler.balance_near_shore,
                balance_vessel_type=config.sampler.balance_vessel_type,
                balance_fishing_type=config.sampler.balance_fishing_type,
                balance_per_scene=config.sampler.balance_per_scene,
                balance_location=config.sampler.balance_location,
                channels_last=config.torch.channels_last,
            )
        else:
            train_ds, train_sampler = dataset.get_random_crop_training_dataset(
                train_df,
                train_image_size=config.train.train_image_size,
                input_channels=config.dataset.channels,
                box_coder=self.box_coder, 
                num_samples=config.sampler.num_samples,
                crop_around_ship_p=config.sampler.crop_around_ship_p,
                individual_augmentations=individual_augmentations,
                augmentations=augmentations,
                normalization=normalization,
                balance_near_shore=config.sampler.balance_near_shore,
                balance_crowd=config.sampler.balance_crowd,
                balance_location=config.sampler.balance_location,
                balance_folder=config.sampler.balance_folder,
                channels_last=config.torch.channels_last,
                shore_root=shore_root,
            )

        valid_ds = dataset.get_validation_dataset(
            valid_df,
            valid_crop_size=config.train.valid_image_size,
            input_channels=config.dataset.channels,
            box_coder=self.box_coder,
            normalization=normalization,
            channels_last=config.torch.channels_last,
        )

        self.master_print(
            "Train",
            "instances",
            len(train_df),
            "dataset",
            len(train_ds),
            "sampler",
            len(train_sampler) if train_sampler else "N/A",
        )
        self.master_print("Valid", "instances", len(valid_df), "dataset", len(valid_ds))

        return train_ds, valid_ds, train_sampler, []

    def build_model(self, config: Dict, loaders: Dict) -> nn.Module:
        model = super().build_model(config)
        logger.debug(f"Train loader sample: {next(iter(loaders['train']))}")
        logger.debug(f"Valid loader sample: {next(iter(loaders['valid']))}")
        nas = ProgressiveNAS(model, config.nas.search_space)
        model = nas.search(config.nas.num_iterations, loaders['train'], loaders['valid'])

        model.yolo_head = YOLOHead(num_classes=config.model.num_classes, anchors=config.model.anchors)
        model.classifier = ConvolutionalClassifier(num_classes=config.model.num_classes)
        return model

    def build_criterions(self, config: Dict) -> Tuple[Dict[str, nn.Module], List[Callback]]:
        criterions, callbacks = super().build_criterions(config)
        
        criterions['yolo_loss'] = YOLOLoss(config.model.num_classes, config.model.anchors)
        criterions['transformer_loss'] = nn.CrossEntropyLoss()
        
        return criterions, callbacks

    def build_metrics(self, config, loaders, model) -> List[Callback]:
        callbacks = super().build_metrics(config, loaders, model)
        
        # Add YOLO and transformer metrics
        callbacks.append(YOLOMetricsCallback(prefix="yolo_metrics"))
        callbacks.append(TransformerMetricsCallback(prefix="transformer_metrics"))
        
        return callbacks
    def build_distributed_params(self, config: DictConfig) -> Dict:
        fp16 = config.optimizer.get("fp16", False)

        if self.distributed:
            local_rank = torch.distributed.get_rank()
            distributed_params = {"rank": local_rank, "syncbn": True}
            if fp16:
                distributed_params["amp"] = True
            if self.cfg.find_unused:
                distributed_params["find_unused_parameters"] = True
        else:
            if fp16:
                distributed_params = {}
                distributed_params["amp"] = True
            else:
                distributed_params = False

        return distributed_params

    def build_optimizer(self, model: nn.Module, loaders: OrderedDict) -> Tuple[Optimizer, Any, List[Callback]]:
        optimizer_config = self.cfg["optimizer"]
        optimizer_use_fp16 = optimizer_config.get("fp16", False)

        optimizer_name = str(optimizer_config["name"])

        wd_on_bias = bool(optimizer_config.get("wd_on_bias", True))
        accumulation_steps = int(optimizer_config.get("accumulation", 1))
        optimizer_params = optimizer_config["params"]
        optimizer_params = scale_learning_rate_for_ddp(optimizer_params)

        optimizer = self.get_optimizer(
            model=model,
            optimizer_name=optimizer_name,
            optimizer_params=optimizer_params,
            apply_weight_decay_to_bias=wd_on_bias,
            layerwise_params=optimizer_config.get("layerwise_params", None),
        )

        if optimizer_use_fp16:
            opt_callback = AMPOptimizerCallback(accumulation_steps=accumulation_steps)
        else:
            opt_callback = OptimizerCallback(accumulation_steps=accumulation_steps, decouple_weight_decay=False)

        scheduler_params = self.cfg["scheduler"]
        num_epochs = self.cfg["train"]["epochs"]

        scheduler_name = str(scheduler_params["scheduler_name"])
        scheduler = get_scheduler(
            optimizer,
            learning_rate=optimizer_params["lr"],
            num_epochs=num_epochs,
            batches_in_epoch=len(loaders["train"]),
            **scheduler_params,
        )

        callbacks = [opt_callback]
        if isinstance(scheduler, (CyclicLR, OneCycleLRWithWarmup)):
            callbacks += [SchedulerCallback(mode="batch")]
        else:
            callbacks += [SchedulerCallback(mode="epoch")]

        self.master_print("Optimizer        :", optimizer_name)
        self.master_print("Model            :", self.cfg.model.config.slug)
        self.master_print(
            "  Parameters     :",
            count_parameters(
                model, ["backbone", "rpn", "roi_heads", "encoder", "decoder", "head", "fuse", "extra_stages", "center", "mask"]
            ),
        )
        self.master_print("  FP16           :", optimizer_use_fp16)
        self.master_print("  Learning rate  :", optimizer_params["lr"])
        self.master_print("  Weight decay   :", optimizer_params.get("weight_decay", 0))
        self.master_print("  WD on bias     :", wd_on_bias)
        self.master_print("  Accumulation   :", accumulation_steps)
        self.master_print("Params           :")
        for k, v in optimizer_params.items():
            self.master_print(f"  {k}:", v)
        self.master_print("Scheduler        :", scheduler_name)

        return optimizer, scheduler, callbacks

    def get_optimizer(
        self,
        model: nn.Module,
        optimizer_name: str,
        optimizer_params: Dict[str, Any],
        apply_weight_decay_to_bias: bool = True,
        layerwise_params=None,
    ) -> Optimizer:
        """
        Construct an Optimizer for given model
        Args:
            model: Model to optimize. Only parameters that require_grad will be used
            optimizer_name: Name of the optimizer (case-insensitive). Supports native pytorch optimizers, apex and
                optimizers from pytorch-optimizers package.
            optimizer_params: Dict of optimizer params (lr, weight_decay, eps, etc)
            apply_weight_decay_to_bias: Whether to apply weight decay on bias parameters. Default is True
        Returns:
            Optimizer instance
        """

        # Optimizer parameter groups
        if layerwise_params is not None:
            if not apply_weight_decay_to_bias:
                raise ValueError("Layerwise params and no wd on bias are mutually exclusive")

            parameters = get_lr_decay_parameters(model, optimizer_params["lr"], layerwise_params)
        else:
            if apply_weight_decay_to_bias:
                parameters = get_optimizable_parameters(model)
            else:
                default_pg, biases_pg = [], []

                for k, v in model.named_parameters():
                    if v.requires_grad:
                        if str.endswith(k, ".bias"):
                            biases_pg.append(v)  # biases
                        else:
                            default_pg.append(v)  # all else

                if apply_weight_decay_to_bias:
                    parameters = default_pg + biases_pg
                else:
                    parameters = default_pg

        optimizer_cls = get_optimizer_cls(optimizer_name)
        optimizer: Optimizer = optimizer_cls(
            parameters,
            **optimizer_params,
        )

        if not apply_weight_decay_to_bias:
            optimizer.add_param_group({"params": biases_pg, "weight_decay": 0.0})

        return optimizer

    def get_experiment_dir(self, config) -> str:
        if self.cfg.experiment:
            return self.cfg.experiment

        current_time = datetime.now().strftime("%y%m%d_%H_%M")
        experiment_slug = f"{current_time}_{config.model.config.slug}_{config.dataset.slug}_{config.loss.slug}_{config.augs.slug}"

        fold = config["dataset"].get("fold", None)
        if fold is not None:
            experiment_slug += f"_fold{fold}"

        log_dir = os.path.join("runs", experiment_slug)
        return log_dir
    
    def get_model(self, model_config: DictConfig) -> nn.Module:
        model: nn.Module = get_detection_model(model_config)
        print(f"[DEBUG] model.box_coder after initialization: {model.box_coder}")
        if model.box_coder is None:
            raise ValueError("model.box_coder is None after model initialization.")
        self.box_coder = model.box_coder
        self.master_print("BoxCoder")
        self.master_print(self.box_coder)
        return model

    def build_loaders(self, config):
        train_ds, valid_ds, train_sampler, dataset_callbacks = self.build_datasets(config)
        
        train_loader = DataLoader(
            train_ds,
            batch_size=config.get('batch_size', 8),
            num_workers=config.get('num_workers', 4),
            sampler=train_sampler,
            pin_memory=True,
            collate_fn=train_ds.get_collate_fn(),
        )
        
        valid_loader = DataLoader(
            valid_ds,
            batch_size=config.get('batch_size', 4),
            num_workers=config.get('num_workers', 4),
            pin_memory=True,
            collate_fn=valid_ds.get_collate_fn(),
        )
        
        return {"train": train_loader, "valid": valid_loader}, dataset_callbacks



    def train(self):
        config = self.cfg
        self.loaders, dataset_callbacks = self.build_loaders(config)
        model = self.build_model(config, self.loaders)
        self.box_coder = model.box_coder  # Set `self.box_coder` after building the model
        # ... rest of your code ...
        #self.loaders, dataset_callbacks = self.build_loaders(config)
        #model = self.build_model(config, self.loaders)
        optimizer, scheduler, optimizer_callbacks = self.build_optimizer(model, self.loaders)
        criterions, criterions_callbacks = self.build_criterions(config)
        metric_callbacks = self.build_metrics(config, self.loaders, model)
        #loaders, dataset_callbacks = self.build_loaders(config)
        nas = ProgressiveNAS(model, config.nas.search_space)
        model = nas.search(config.nas.num_iterations, self.loaders['train'], self.loaders['valid'])
        #optimizer, scheduler, optimizer_callbacks = self.build_optimizer(model, loaders)
        #criterions, criterions_callbacks = self.build_criterions(config)
        #metric_callbacks = self.build_metrics(config, loaders, model)

        experiment_dir = self.get_experiment_dir(config)
        self.experiment_dir = experiment_dir

        if self.is_master:
            os.makedirs(experiment_dir, exist_ok=True)
            dst_config_fname = os.path.join(experiment_dir, "config.yaml")
            with open(dst_config_fname, "w") as f:
                OmegaConf.save(self.cfg, f)

        # model training
        runner_config = config.get("runner", {})
        input_key = runner_config.get("input_key", INPUT_SCENE_ID_KEY)
        output_key = runner_config.get("output_key", None)

        if input_key is not None and not isinstance(input_key, str):
            input_key = OmegaConf.to_container(input_key, resolve=True)
        if output_key is not None and not isinstance(output_key, str):
            output_key = OmegaConf.to_container(output_key, resolve=True)

        device = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'
        runner = SupervisedRunner(input_key=input_key, output_key=output_key, device=device)

        extra_callbacks = [TimerCallback()]
        # try:
        runner.train(
            fp16=self.build_distributed_params(config),
            model=model,
            criterion=criterions,
            optimizer=optimizer,
            scheduler=scheduler,
            callbacks=dataset_callbacks + metric_callbacks + optimizer_callbacks + criterions_callbacks + extra_callbacks,
            loaders=loaders,
            logdir=experiment_dir,
            num_epochs=config.train.epochs,
            verbose=True,
            main_metric=config.runner.main_metric,
            minimize_metric=config.runner.main_metric_minimize,
            checkpoint_data=self.get_checkpoint_data(),
        )
        # except Exception as e:
        #     with open(os.path.join(experiment_dir, "exception.log"), "w") as logf:
        #         logf.write(str(e))
        #         traceback.print_exc(file=logf)
        #         print(e)

        self.on_experiment_finished()

    def get_checkpoint_data(self):
        return {
            "config": OmegaConf.to_container(self.cfg, resolve=True),
            "version": get_version(),
        }

    def on_experiment_finished(self):
        pass

    def build_quircks(self, config: DictConfig) -> List[Callback]:
        extra_callbacks = []

        if config.ema.callback is not None:
            cb = hydra.utils.instantiate(config.ema.callback)
            extra_callbacks.append(cb)
            self.master_print("Using EMA Callback", cb)

        return extra_callbacks


class TrainOnlyCriterionCallback(CriterionCallback):
    def on_batch_end(self, runner) -> None:
        if runner.is_train_loader:
            return super(TrainOnlyCriterionCallback, self).on_batch_end(runner)
        else:
            runner.batch_metrics[self.prefix] = torch.tensor(0, device="mps")



@hydra.main(config_path="configs", config_name="multilabel_circle_net", version_base=None)
def main(cfg: DictConfig) -> None:
    torch.mps.empty_cache()
    torch.manual_seed(int(cfg.seed))

    torch.set_anomaly_enabled(cfg.torch.detect_anomaly)

    torch.backends.cudnn.deterministic = cfg.torch.deterministic
    torch.backends.cudnn.benchmark = cfg.torch.benchmark

    # The flag below controls whether to allow TF32 on matmul. This flag defaults to True.
    torch.backends.cuda.matmul.allow_tf32 = cfg.torch.cuda_allow_tf32
    # The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
    torch.backends.cudnn.allow_tf32 = cfg.torch.cudnn.allow_tf32

    if cfg.dataset.data_dir is None:
        raise ValueError("--data-dir must be set")

    CenternetPipeline(cfg).train()


if __name__ == "__main__":
    main()


#then this file centernet_unet.py

from functools import partial
from typing import List, Dict
from omegaconf import DictConfig
from pytorch_toolbelt.modules import (
    ACT_RELU,
    decoders as D,
    EncoderModule,
    ResidualDeconvolutionUpsample2d,
)
from torch import nn
import torch
import timm
from hydra.utils import instantiate
import random
import torch.nn.functional as F
from xview3.centernet.models.custom_encoder import CustomTimmEncoder
from .unet_blocks import IRBlock, get_unet_block, SqueezeExcitation
from .necks import FPNNeck
from xview3.nas import NAS as NeuralArchitectureSearch
from .heads.yolo_head import YOLOHead
from .transformer_classifier import ConvolutionalClassifier
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,
)
__all__ = ["CenterNetUNetModel"]

class CenterNetUNetModel(nn.Module):
    def __init__(
        self,
        encoder: EncoderModule,
        decoder_features: List[int],
        head: nn.Module,
        box_coder,
        neck_out_channels: int = 256,  # Add this parameter
        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)],
        **kwargs
    ):
        super().__init__()
        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)]
        )
        upsample_block = {
            "nearest": nn.UpsamplingNearest2d,
            "bilinear": nn.UpsamplingBilinear2d,
            "rdtsc": ResidualDeconvolutionUpsample2d,
            "shuffle": nn.PixelShuffle,
        }[upsample_type]

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

        #self.neck = FPNNeck(self.decoder_channels, neck_out_channels)
        #self.neck_conv = nn.Conv2d(self.neck_output_channels, neck_out_channels, kernel_size=1)

        self.head = head
        self.head_output_channels = head.out_channels  # Ensure your head module defines 'out_channels'

        # Ensure self.nas is a module that can be called during forward
        self.nas = NeuralArchitectureSearch(
            search_space={
                'num_layers': (2, 5),
                'num_filters': (32, 128),
                'kernel_size': (3, 5),
                'activation': ['relu', 'leaky_relu']
            },
            evaluation_function=self.evaluate_architecture
        )
        self.nas_output_channels = self.nas.output_channels  # Define this in your NAS module

        self.yolo_head = YOLOHead(num_classes=num_classes, anchors=anchors)
        self.yolo_output_channels = self.yolo_head.num_anchors * (num_classes + 5)  # Adjust as per YOLO implementation

        self.classifier = ConvolutionalClassifier(in_channels=neck_out_channels, num_classes=num_classes)
        self.classifier_output_channels = self.classifier.out_channels
        self.box_coder = box_coder
        self.head_output_channels = getattr(self.head, 'out_channels', 256)
        self.nas_output_channels = getattr(self.nas, 'output_channels', 256)
        self.yolo_output_channels = self.yolo_head.num_anchors * (num_classes + 5)
        self.classifier_output_channels = getattr(self.classifier, 'out_channels', num_classes)

        # Ensure all channel values are integers
        total_input_channels = (
            int(self.head_output_channels) +
            int(self.nas_output_channels) +
            int(self.yolo_output_channels) +
            int(self.classifier_output_channels)
        )
        self.fusion = nn.Conv2d(total_input_channels, self.head_output_channels, kernel_size=1)

        self.dropout = nn.Dropout(p=0.5)
        #self.batch_norm = FilterResponseNormalization(256)
        #self.batch_norm = FilterResponseNormalization(neck_out_channels * len(self.decoder_channels))
        # Corrected code
        self.decoder_channels = decoder_features
        #self.num_decoded_features = len(encoder_channels) - 1  # The number of decoder stages
        self.num_decoded_features = len(self.decoder_channels)
        #self.num_decoder_features_used_by_neck = self.num_decoded_features
        #self.neck_output_channels = neck_out_channels * self.num_decoder_features_used_by_neck
        self.neck_output_channels = neck_out_channels * self.num_decoded_features
        print(f"[DEBUG] neck_output_channels: {self.neck_output_channels}")
        num_decoder_stages = self.num_decoded_features
        if len(self.decoder_channels) < num_decoder_stages:
            # Extend decoder_channels by repeating the last value
            self.decoder_channels.extend([self.decoder_channels[-1]] * (num_decoder_stages - len(self.decoder_channels)))
        elif len(self.decoder_channels) > num_decoder_stages:
            # Trim decoder_channels to the required length
            self.decoder_channels = self.decoder_channels[:num_decoder_stages]
        self.decoder = D.UNetDecoder(
            encoder_channels,
            self.decoder_channels,
            unet_block=unet_block,
            upsample_block=upsample_block,
        )
        self.neck = FPNNeck(self.decoder_channels, neck_out_channels)
        self.neck_conv = nn.Conv2d(
            in_channels=self.neck_output_channels,
            out_channels=neck_out_channels,
            kernel_size=1
        )

        print(f"[DEBUG] neck_output_channels: {self.neck_output_channels}")

        self.batch_norm = FilterResponseNormalization(self.neck_output_channels)

    def evaluate_architecture(self, architecture):
        # Implement a simple evaluation function
        return random.random()  # Replace with actual evaluation logic

    def forward(self, x):
        x = self.initial_upsample(x)
        print(f"[DEBUG] x shape after initial_upsample: {x.shape}")
        feature_maps = self.encoder(x)
        print(f"[DEBUG] feature_maps type: {type(feature_maps)}")
        if isinstance(feature_maps, list):
            for idx, fmap in enumerate(feature_maps):
                print(f"[DEBUG] feature_maps[{idx}] shape: {fmap.shape}")
        else:
            print("[ERROR] feature_maps is not a list")
            print(f"feature_maps shape: {feature_maps.shape}")
            print(f"feature_maps type: {type(feature_maps)}")
            # Convert feature_maps to a list if needed
            feature_maps = [feature_maps]
        for extra_block in self.extra_encoder_blocks:
            feature_maps.append(extra_block(feature_maps[-1]))
        decoded_features = self.decoder(feature_maps)
        print(f"[DEBUG] decoded_features type: {type(decoded_features)}")
        if isinstance(decoded_features, list):
            for idx, fmap in enumerate(decoded_features):
                print(f"[DEBUG] decoded_features[{idx}] shape: {fmap.shape}")
        else:
            print("[ERROR] decoded_features is not a list")
            print(f"decoded_features shape: {decoded_features.shape}")
            print(f"decoded_features type: {type(decoded_features)}")
            # Convert decoded_features to a list if needed
            decoded_features = [decoded_features]
        neck_features = self.neck(decoded_features)
        print(f"[DEBUG] neck_features shape before batch_norm: {neck_features.shape}")
        print(f"[DEBUG] neck_features channels: {neck_features.size(1)}")
        print(f"[DEBUG] neck_features shape before batch_norm: {neck_features.shape}")
        neck_features = self.batch_norm(neck_features)
        print(f"[DEBUG] neck_features shape after batch_norm: {neck_features.shape}")
        neck_features = self.dropout(neck_features)
        neck_features = self.neck_conv(neck_features)

        #self.head.channels = self.nas_output_channels  # Assuming self.nas_output_channels is the correct channel size
        head_output = self.head(neck_features)
        nas_output = self.nas(neck_features)
        yolo_output = self.yolo_head(neck_features)
        classification_output = self.classifier(neck_features)

        # Fusion layer
        fused_output = self.fusion(torch.cat([head_output, nas_output, yolo_output, classification_output], dim=1))

        return {
            CENTERNET_OUTPUT_OBJECTNESS_MAP: fused_output[:, :1],
            CENTERNET_OUTPUT_VESSEL_MAP: fused_output[:, 1:2],
            CENTERNET_OUTPUT_FISHING_MAP: fused_output[:, 2:3],
            CENTERNET_OUTPUT_SIZE: fused_output[:, 3:4],
            CENTERNET_OUTPUT_OFFSET: fused_output[:, 4:6],
        }

    @classmethod
    def from_config(cls, **config):
        # Extract and pop the encoder configuration
        encoder_config = config.pop('encoder')
        num_channels = config.pop('num_channels', 3)
        model_name = encoder_config.get('model_name', 'tf_efficientnet_b5_ns')
        pretrained = encoder_config.get('pretrained', False)

        # Instantiate the custom encoder
        encoder_module = CustomTimmEncoder(
            model_name=model_name,
            num_channels=num_channels,
            features_only=True,
            pretrained=pretrained
        )

        # Pop other configurations
        input_upsample_factor = config.pop("input_upsample_factor", 1)
        activation = config.pop('activation', ACT_RELU)
        num_extra_blocks = config.pop('num_extra_blocks', 0)
        num_classes = config.pop("num_classes", 2)
        anchors = config.pop("anchors", [(10, 13), (16, 30), (33, 23)])

        # Pop decoder configuration
        decoder_config = config.pop('decoder')
        decoder_features = decoder_config['channels']
        block_type = decoder_config.get("block_type", "UnetBlock")
        upsample_type = decoder_config.get("upsample_type", "bilinear")

        neck_out_channels = config.get('neck_out_channels', 256)
        neck_output_channels = neck_out_channels * len(decoder_features)
        print(f"Neck output channels: {neck_output_channels}")

        # Instantiate head and box_coder
        head_config = config.pop('head')
        head = instantiate(head_config, channels=neck_output_channels)

        box_coder_config = config.pop('box_coder')
        box_coder = instantiate(
            box_coder_config,
            output_stride=encoder_module.strides[0] // (head.upsample_factor * input_upsample_factor),
        )

        # Now initialize the model instance with the remaining config
        return cls(
            encoder=encoder_module,
            decoder_features=decoder_features,
            head=head,
            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,
            **config
        )

def adjust_encoder_input_channels(encoder_module, in_chans):
    # Access the underlying encoder model
    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
        print(f"[DEBUG] Before adjustment, encoder.conv_stem.in_channels: {old_conv.in_channels}")

        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
        )

        # Initialize new_conv weights
        with torch.no_grad():
            # Copy weights for existing channels
            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:
                # Initialize added 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, in_channels, num_classes):
        super().__init__()
        self.backbone = CenterNetUNetModel(in_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)

#then this file train_multilabel.py

import sys
import traceback
import warnings
from typing import Dict, Any, OrderedDict, Tuple, Optional, List
import catalyst
import os
import numpy as np
import pandas as pd
import torch
import hydra
from torch.optim import Optimizer
from torch.utils.data import Dataset, Sampler, DataLoader
from torch import nn
from omegaconf import DictConfig
from functools import partial
import albumentations as A
from catalyst.callbacks import ControlFlowCallback, CriterionCallback, SchedulerCallback, TimerCallback
from catalyst.core import Callback
from pytorch_toolbelt.utils.catalyst import (
    ShowPolarBatchesCallback,
    BestMetricCheckpointCallback,
    HyperParametersCallback,
    MixupCriterionCallback,
    MixupInputCallback,
)
from pytorch_toolbelt.optimization.functional import freeze_model, get_lr_decay_parameters, get_optimizable_parameters
from pytorch_toolbelt.utils import fs, transfer_weights, count_parameters
from pytorch_toolbelt.utils.distributed import is_main_process
from xview3.centernet.bboxer import MultilabelCircleNetCoder
from xview3.early_stopping import EarlyStoppingCallback
from xview3.nas import ProgressiveNAS
from xview3.centernet.models.heads.yolo_head import YOLOHead
from xview3.centernet.models.transformer_classifier import ConvolutionalClassifier
from xview3.centernet.losses.circlenet_loss import CircleNetLoss, YOLOLoss
from xview3.centernet import *
from xview3.pipeline import *
from xview3.factory import *
from xview3.augmentations import *
from xview3.samplers import *
from xview3.dataset import *
from xview3.ema import *
from xview3.evaluation import *
from xview3.inference import *
from xview3.metric import *
from xview3.constants import INPUT_SCENE_ID_KEY
import logging

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

def warn_with_traceback(message, category, filename, lineno, file=None, line=None):
    log = file if hasattr(file, "write") else sys.stderr
    traceback.print_stack(file=log)
    log.write(warnings.formatwarning(message, category, filename, lineno))

warnings.showwarning = warn_with_traceback

if 'XVIEW3_DIR' not in os.environ:
    os.environ['XVIEW3_DIR'] = '/Volumes/Crucial/xview3_dataset'

class CenternetPipeline(Pipeline):
    def __init__(self, cfg):
        super().__init__(cfg)
        self.box_coder: Optional[MultilabelCircleNetCoder] = None
        self.shore_root = os.path.join(os.environ.get('XVIEW3_DIR', ''), "full_dataset", "shoreline")

    def build_datasets(self, config) -> Tuple[Dataset, Dataset, Optional[Sampler], List[Callback]]:
        dataset = MultilabelCircleNetDataModule(data_dir=os.path.join(os.environ.get('XVIEW3_DIR', '/Volumes/Crucial/xview3_dataset'), 'full_dataset'))
        self.box_coder = MultilabelCircleNetCoder(
            image_size=(config.train.train_image_size[0], config.train.train_image_size[1]),
            output_stride=4,  # Adjust this value as needed
            max_objects=512,  # Adjust this value as needed
            heatmap_encoding="umich",
            labels_encoding="circle",
            ignore_value=-1
        )
        print(f"[DEBUG] self.box_coder in build_datasets: {self.box_coder}")
        if self.box_coder is None:
            raise ValueError("self.box_coder is None in build_datasets.")

        train_df, valid_df, _, shore_root = dataset.train_val_split(
            splitter=config.dataset.splitter,
            fold=config.dataset.fold,
            num_folds=config.dataset.num_folds,
        )

        self.valid_df = valid_df.copy()  # Keep a copy of unmodified groundtruth

        self.master_print("Dataset", config.dataset.slug)
        self.master_print("\tignore_low_confidence_detections", config.dataset.ignore_low_confidence_detections)
        self.master_print("\tignore_low_confidence_labels    ", config.dataset.ignore_low_confidence_labels)
        self.master_print("\tfilter_low_confidence_objects   ", config.dataset.filter_low_confidence_objects)

        self.master_print("Train dataset")
        self.master_print("\tUnknown vessels", np.isnan(train_df.is_vessel.values.astype(float)).sum())
        self.master_print("\tUnknown fishing", np.isnan(train_df.is_fishing.values.astype(float)).sum())
        self.master_print("Valid dataset")
        self.master_print("\tUnknown vessels", np.isnan(valid_df.is_vessel.values.astype(float)).sum())
        self.master_print("\tUnknown fishing", np.isnan(valid_df.is_fishing.values.astype(float)).sum())

        if config.dataset.filter_low_confidence_objects and config.dataset.ignore_low_confidence_labels:
            raise ValueError(f"filter_low_confidence_objects and ignore_low_confidence_labels are mutually exclusive")

        if config.dataset.filter_low_confidence_objects:
            train_df = filter_low_confidence_objects(train_df)
            valid_df = filter_low_confidence_objects(valid_df)
            self.master_print("Excluding low-confidence objects from training & validation")
        elif config.dataset.ignore_low_confidence_labels:
            train_df = ignore_low_confidence_objects(train_df)
            valid_df = ignore_low_confidence_objects(valid_df)
            self.master_print("Ignoring low-confidence labels from training & validation")

        self.shore_root = os.path.join(os.environ.get('XVIEW3_DIR', ''), "full_dataset", "shoreline")

        normalization = build_normalization(config.normalization)

        augmentations = [
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.5),
            A.RandomRotate90(p=0.5),
            A.Resize(height=config.train.train_image_size[0], width=config.train.train_image_size[1], always_apply=True),
        ]

        transform = A.Compose(
            augmentations,
            keypoint_params=A.KeypointParams(
                format='xys',
                label_fields=['confidences', 'is_vessel', 'is_fishing', 'is_near_shore'],
                remove_invisible=False
            )
        )


        individual_augmentations = dict(
            [(key, A.Compose(build_augmentations([value]))) for key, value in config.augs.individual.items() if value is not None]
        )

        for channel in config.dataset.channels:
            if channel not in normalization:
                normalization[channel] = A.NoOp()

        if config.sampler.sampler_type == "each_object":
            train_ds, train_sampler = dataset.get_random_crop_each_object_dataset(
                train_df,
                train_image_size=config.train.train_image_size,
                input_channels=config.dataset.channels,
                box_coder=self.box_coder,  # Pass `self.box_coder` here
                num_samples=config.sampler.num_samples,
                individual_augmentations=individual_augmentations,
                augmentations=augmentations,
                transform=transform,
                normalization=normalization,
                balance_crowd=config.sampler.balance_crowd,
                balance_near_shore=config.sampler.balance_near_shore,
                balance_vessel_type=config.sampler.balance_vessel_type,
                balance_fishing_type=config.sampler.balance_fishing_type,
                balance_per_scene=config.sampler.balance_per_scene,
                balance_location=config.sampler.balance_location,
                channels_last=config.torch.channels_last,
            )
        else:
            train_ds, train_sampler = dataset.get_random_crop_training_dataset(
                train_df,
                train_image_size=config.train.train_image_size,
                input_channels=config.dataset.channels,
                box_coder=self.box_coder, 
                num_samples=config.sampler.num_samples,
                crop_around_ship_p=config.sampler.crop_around_ship_p,
                individual_augmentations=individual_augmentations,
                augmentations=augmentations,
                normalization=normalization,
                balance_near_shore=config.sampler.balance_near_shore,
                balance_crowd=config.sampler.balance_crowd,
                balance_location=config.sampler.balance_location,
                balance_folder=config.sampler.balance_folder,
                channels_last=config.torch.channels_last,
                shore_root=shore_root,
            )

        valid_ds = dataset.get_validation_dataset(
            valid_df,
            valid_crop_size=config.train.valid_image_size,
            input_channels=config.dataset.channels,
            box_coder=self.box_coder,
            normalization=normalization,
            channels_last=config.torch.channels_last,
        )

        self.master_print(
            "Train",
            "instances",
            len(train_df),
            "dataset",
            len(train_ds),
            "sampler",
            len(train_sampler) if train_sampler else "N/A",
        )
        self.master_print("Valid", "instances", len(valid_df), "dataset", len(valid_ds))

        return train_ds, valid_ds, train_sampler, []

    def build_model(self, config: Dict, loaders: Dict) -> nn.Module:
        model = super().build_model(config)
        logger.debug(f"Train loader sample: {next(iter(loaders['train']))}")
        logger.debug(f"Valid loader sample: {next(iter(loaders['valid']))}")
        nas = ProgressiveNAS(model, config.nas.search_space)
        model = nas.search(config.nas.num_iterations, loaders['train'], loaders['valid'])

        model.yolo_head = YOLOHead(num_classes=config.model.num_classes, anchors=config.model.anchors)
        model.classifier = ConvolutionalClassifier(num_classes=config.model.num_classes)
        return model

    def build_criterions(self, config: Dict) -> Tuple[Dict[str, nn.Module], List[Callback]]:
        criterions, callbacks = super().build_criterions(config)
        
        criterions['yolo_loss'] = YOLOLoss(config.model.num_classes, config.model.anchors)
        criterions['transformer_loss'] = nn.CrossEntropyLoss()
        
        return criterions, callbacks

    def build_metrics(self, config, loaders, model) -> List[Callback]:
        callbacks = super().build_metrics(config, loaders, model)
        
        # Add YOLO and transformer metrics
        callbacks.append(YOLOMetricsCallback(prefix="yolo_metrics"))
        callbacks.append(TransformerMetricsCallback(prefix="transformer_metrics"))
        
        return callbacks
    def build_distributed_params(self, config: DictConfig) -> Dict:
        fp16 = config.optimizer.get("fp16", False)

        if self.distributed:
            local_rank = torch.distributed.get_rank()
            distributed_params = {"rank": local_rank, "syncbn": True}
            if fp16:
                distributed_params["amp"] = True
            if self.cfg.find_unused:
                distributed_params["find_unused_parameters"] = True
        else:
            if fp16:
                distributed_params = {}
                distributed_params["amp"] = True
            else:
                distributed_params = False

        return distributed_params

    def build_optimizer(self, model: nn.Module, loaders: OrderedDict) -> Tuple[Optimizer, Any, List[Callback]]:
        optimizer_config = self.cfg["optimizer"]
        optimizer_use_fp16 = optimizer_config.get("fp16", False)

        optimizer_name = str(optimizer_config["name"])

        wd_on_bias = bool(optimizer_config.get("wd_on_bias", True))
        accumulation_steps = int(optimizer_config.get("accumulation", 1))
        optimizer_params = optimizer_config["params"]
        optimizer_params = scale_learning_rate_for_ddp(optimizer_params)

        optimizer = self.get_optimizer(
            model=model,
            optimizer_name=optimizer_name,
            optimizer_params=optimizer_params,
            apply_weight_decay_to_bias=wd_on_bias,
            layerwise_params=optimizer_config.get("layerwise_params", None),
        )

        if optimizer_use_fp16:
            opt_callback = AMPOptimizerCallback(accumulation_steps=accumulation_steps)
        else:
            opt_callback = OptimizerCallback(accumulation_steps=accumulation_steps, decouple_weight_decay=False)

        scheduler_params = self.cfg["scheduler"]
        num_epochs = self.cfg["train"]["epochs"]

        scheduler_name = str(scheduler_params["scheduler_name"])
        scheduler = get_scheduler(
            optimizer,
            learning_rate=optimizer_params["lr"],
            num_epochs=num_epochs,
            batches_in_epoch=len(loaders["train"]),
            **scheduler_params,
        )

        callbacks = [opt_callback]
        if isinstance(scheduler, (CyclicLR, OneCycleLRWithWarmup)):
            callbacks += [SchedulerCallback(mode="batch")]
        else:
            callbacks += [SchedulerCallback(mode="epoch")]

        self.master_print("Optimizer        :", optimizer_name)
        self.master_print("Model            :", self.cfg.model.config.slug)
        self.master_print(
            "  Parameters     :",
            count_parameters(
                model, ["backbone", "rpn", "roi_heads", "encoder", "decoder", "head", "fuse", "extra_stages", "center", "mask"]
            ),
        )
        self.master_print("  FP16           :", optimizer_use_fp16)
        self.master_print("  Learning rate  :", optimizer_params["lr"])
        self.master_print("  Weight decay   :", optimizer_params.get("weight_decay", 0))
        self.master_print("  WD on bias     :", wd_on_bias)
        self.master_print("  Accumulation   :", accumulation_steps)
        self.master_print("Params           :")
        for k, v in optimizer_params.items():
            self.master_print(f"  {k}:", v)
        self.master_print("Scheduler        :", scheduler_name)

        return optimizer, scheduler, callbacks

    def get_optimizer(
        self,
        model: nn.Module,
        optimizer_name: str,
        optimizer_params: Dict[str, Any],
        apply_weight_decay_to_bias: bool = True,
        layerwise_params=None,
    ) -> Optimizer:
        """
        Construct an Optimizer for given model
        Args:
            model: Model to optimize. Only parameters that require_grad will be used
            optimizer_name: Name of the optimizer (case-insensitive). Supports native pytorch optimizers, apex and
                optimizers from pytorch-optimizers package.
            optimizer_params: Dict of optimizer params (lr, weight_decay, eps, etc)
            apply_weight_decay_to_bias: Whether to apply weight decay on bias parameters. Default is True
        Returns:
            Optimizer instance
        """

        # Optimizer parameter groups
        if layerwise_params is not None:
            if not apply_weight_decay_to_bias:
                raise ValueError("Layerwise params and no wd on bias are mutually exclusive")

            parameters = get_lr_decay_parameters(model, optimizer_params["lr"], layerwise_params)
        else:
            if apply_weight_decay_to_bias:
                parameters = get_optimizable_parameters(model)
            else:
                default_pg, biases_pg = [], []

                for k, v in model.named_parameters():
                    if v.requires_grad:
                        if str.endswith(k, ".bias"):
                            biases_pg.append(v)  # biases
                        else:
                            default_pg.append(v)  # all else

                if apply_weight_decay_to_bias:
                    parameters = default_pg + biases_pg
                else:
                    parameters = default_pg

        optimizer_cls = get_optimizer_cls(optimizer_name)
        optimizer: Optimizer = optimizer_cls(
            parameters,
            **optimizer_params,
        )

        if not apply_weight_decay_to_bias:
            optimizer.add_param_group({"params": biases_pg, "weight_decay": 0.0})

        return optimizer

    def get_experiment_dir(self, config) -> str:
        if self.cfg.experiment:
            return self.cfg.experiment

        current_time = datetime.now().strftime("%y%m%d_%H_%M")
        experiment_slug = f"{current_time}_{config.model.config.slug}_{config.dataset.slug}_{config.loss.slug}_{config.augs.slug}"

        fold = config["dataset"].get("fold", None)
        if fold is not None:
            experiment_slug += f"_fold{fold}"

        log_dir = os.path.join("runs", experiment_slug)
        return log_dir
    
    def get_model(self, model_config: DictConfig) -> nn.Module:
        model: nn.Module = get_detection_model(model_config)
        print(f"[DEBUG] model.box_coder after initialization: {model.box_coder}")
        if model.box_coder is None:
            raise ValueError("model.box_coder is None after model initialization.")
        self.box_coder = model.box_coder
        self.master_print("BoxCoder")
        self.master_print(self.box_coder)
        return model

    def build_loaders(self, config):
        train_ds, valid_ds, train_sampler, dataset_callbacks = self.build_datasets(config)
        
        train_loader = DataLoader(
            train_ds,
            batch_size=config.get('batch_size', 8),
            num_workers=config.get('num_workers', 4),
            sampler=train_sampler,
            pin_memory=True,
            collate_fn=train_ds.get_collate_fn(),
        )
        
        valid_loader = DataLoader(
            valid_ds,
            batch_size=config.get('batch_size', 4),
            num_workers=config.get('num_workers', 4),
            pin_memory=True,
            collate_fn=valid_ds.get_collate_fn(),
        )
        
        return {"train": train_loader, "valid": valid_loader}, dataset_callbacks



    def train(self):
        config = self.cfg
        self.loaders, dataset_callbacks = self.build_loaders(config)
        model = self.build_model(config, self.loaders)
        self.box_coder = model.box_coder  # Set `self.box_coder` after building the model
        # ... rest of your code ...
        #self.loaders, dataset_callbacks = self.build_loaders(config)
        #model = self.build_model(config, self.loaders)
        optimizer, scheduler, optimizer_callbacks = self.build_optimizer(model, self.loaders)
        criterions, criterions_callbacks = self.build_criterions(config)
        metric_callbacks = self.build_metrics(config, self.loaders, model)
        #loaders, dataset_callbacks = self.build_loaders(config)
        nas = ProgressiveNAS(model, config.nas.search_space)
        model = nas.search(config.nas.num_iterations, self.loaders['train'], self.loaders['valid'])
        #optimizer, scheduler, optimizer_callbacks = self.build_optimizer(model, loaders)
        #criterions, criterions_callbacks = self.build_criterions(config)
        #metric_callbacks = self.build_metrics(config, loaders, model)

        experiment_dir = self.get_experiment_dir(config)
        self.experiment_dir = experiment_dir

        if self.is_master:
            os.makedirs(experiment_dir, exist_ok=True)
            dst_config_fname = os.path.join(experiment_dir, "config.yaml")
            with open(dst_config_fname, "w") as f:
                OmegaConf.save(self.cfg, f)

        # model training
        runner_config = config.get("runner", {})
        input_key = runner_config.get("input_key", INPUT_SCENE_ID_KEY)
        output_key = runner_config.get("output_key", None)

        if input_key is not None and not isinstance(input_key, str):
            input_key = OmegaConf.to_container(input_key, resolve=True)
        if output_key is not None and not isinstance(output_key, str):
            output_key = OmegaConf.to_container(output_key, resolve=True)

        device = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'
        runner = SupervisedRunner(input_key=input_key, output_key=output_key, device=device)

        extra_callbacks = [TimerCallback()]
        # try:
        runner.train(
            fp16=self.build_distributed_params(config),
            model=model,
            criterion=criterions,
            optimizer=optimizer,
            scheduler=scheduler,
            callbacks=dataset_callbacks + metric_callbacks + optimizer_callbacks + criterions_callbacks + extra_callbacks,
            loaders=loaders,
            logdir=experiment_dir,
            num_epochs=config.train.epochs,
            verbose=True,
            main_metric=config.runner.main_metric,
            minimize_metric=config.runner.main_metric_minimize,
            checkpoint_data=self.get_checkpoint_data(),
        )
        # except Exception as e:
        #     with open(os.path.join(experiment_dir, "exception.log"), "w") as logf:
        #         logf.write(str(e))
        #         traceback.print_exc(file=logf)
        #         print(e)

        self.on_experiment_finished()

    def get_checkpoint_data(self):
        return {
            "config": OmegaConf.to_container(self.cfg, resolve=True),
            "version": get_version(),
        }

    def on_experiment_finished(self):
        pass

    def build_quircks(self, config: DictConfig) -> List[Callback]:
        extra_callbacks = []

        if config.ema.callback is not None:
            cb = hydra.utils.instantiate(config.ema.callback)
            extra_callbacks.append(cb)
            self.master_print("Using EMA Callback", cb)

        return extra_callbacks


class TrainOnlyCriterionCallback(CriterionCallback):
    def on_batch_end(self, runner) -> None:
        if runner.is_train_loader:
            return super(TrainOnlyCriterionCallback, self).on_batch_end(runner)
        else:
            runner.batch_metrics[self.prefix] = torch.tensor(0, device="mps")



@hydra.main(config_path="configs", config_name="multilabel_circle_net", version_base=None)
def main(cfg: DictConfig) -> None:
    torch.mps.empty_cache()
    torch.manual_seed(int(cfg.seed))

    torch.set_anomaly_enabled(cfg.torch.detect_anomaly)

    torch.backends.cudnn.deterministic = cfg.torch.deterministic
    torch.backends.cudnn.benchmark = cfg.torch.benchmark

    # The flag below controls whether to allow TF32 on matmul. This flag defaults to True.
    torch.backends.cuda.matmul.allow_tf32 = cfg.torch.cuda_allow_tf32
    # The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
    torch.backends.cudnn.allow_tf32 = cfg.torch.cudnn.allow_tf32

    if cfg.dataset.data_dir is None:
        raise ValueError("--data-dir must be set")

    CenternetPipeline(cfg).train()


if __name__ == "__main__":
    main()


#next file necks.py

import collections
from typing import List, Dict
from pytorch_toolbelt.modules import FPNFuse
from torch import nn, Tensor

__all__ = ["TensorAtIndex", "PickTensors", "HyperColumnNeck"]


class TensorAtIndex(nn.Module):
    def __init__(self, channels: List[int], index: int = 0):
        super().__init__()
        self.index = index
        self.channels = channels[index]

    def forward(self, feature_maps: List[Tensor]) -> Tensor:
        return feature_maps[self.index]


class PickTensors(nn.Module):
    def __init__(self, channels: List[int], indexes: Dict[int, str]):
        super().__init__()
        self.indexes = indexes
        self.channels = channels[0]

    def forward(self, feature_maps: List[Tensor]) -> Dict[str, Tensor]:
        return collections.OrderedDict((name, feature_maps[index]) for index, name in self.indexes.items())


class HyperColumnNeck(nn.Module):
    def __init__(self, channels: List[int], dropout_rate=0.0):
        super().__init__()
        self.fuse = FPNFuse(align_corners=True)
        self.channels = sum(channels)
        self.dropout = nn.Dropout2d(p=dropout_rate, inplace=True)

    def forward(self, feature_maps: List[Tensor]) -> Tensor:
        return self.dropout(self.fuse(feature_maps))

class FPNNeck(nn.Module):
    def __init__(self, channels: List[int], out_channels: int):
        super().__init__()
        self.fpn = FPNFuse(channels, out_channels, mode='bilinear', align_corners=False)
        self.out_conv = nn.Conv2d(out_channels * len(channels), out_channels, kernel_size=1)

    def forward(self, feature_maps: List[Tensor]) -> Tensor:
        fpn_output = self.fpn(feature_maps)
        return self.out_conv(fpn_output)

#and now file yolo_head.py

from torch import nn, Tensor
import torch
from pytorch_toolbelt.modules import FPNFuse
from typing import List, Tuple, Union

class AnchorFreeHead(nn.Module):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.cls_conv = nn.Conv2d(in_channels, num_classes, 1)
        self.reg_conv = nn.Conv2d(in_channels, 4, 1)
        self.obj_conv = nn.Conv2d(in_channels, 1, 1)

    def forward(self, x):
        cls_pred = self.cls_conv(x)
        reg_pred = self.reg_conv(x)
        obj_pred = self.obj_conv(x)
        return torch.cat([reg_pred, obj_pred, cls_pred], dim=1)

class YOLOHead(nn.Module):
    def __init__(self, num_classes, anchors):
        super(YOLOHead, self).__init__()
        self.num_classes = num_classes
        self.anchors = anchors
        self.num_anchors = len(anchors)
        self.fpn = FPNFuse([256, 512, 1024], 256)
        self.out_channels = self.num_anchors * (5 + self.num_classes)
        self.conv = nn.Conv2d(256, self.out_channels, 1)

    def forward(self, x):
        print(f"[DEBUG] x shape in YOLOHead before fpn: {x.shape}")
        x = self.fpn(x)
        print(f"[DEBUG] x shape in YOLOHead after fpn: {x.shape}")
        batch_size, _, height, width = x.shape
        # Continue with the rest of the code
        return self.conv(x).view(batch_size, len(self.anchors), 5 + self.num_classes, height, width).permute(0, 1, 3, 4, 2)


#lastely file decoupled_head.py

from collections import OrderedDict
from typing import Dict, Tuple, List

from pytorch_toolbelt.modules import ACT_RELU, get_activation_block, conv1x1
from torch import nn, Tensor

from ...constants import (
    CENTERNET_OUTPUT_VESSEL_MAP,
    CENTERNET_OUTPUT_FISHING_MAP,
    CENTERNET_OUTPUT_OBJECTNESS_MAP,
    CENTERNET_OUTPUT_SIZE,
    CENTERNET_OUTPUT_OFFSET,
)

__all__ = ["DecoupledHeadGroupNormLateShuffle"]


class FeatureRecombination(nn.Module):
    def __init__(self, channels, modulating_features):
        super().__init__()
        self.project = conv1x1(modulating_features, channels)

    def forward(self, x, modulation):
        return x * self.project(modulation).sigmoid()


def build_regression_block(
    input_channels: int,
    embedding_size: int,
    dropout_rate: float,
    num_blocks: int,
    activation: str,
    extra_final_layers: List[Tuple[str, nn.Module]],
    num_groups=16,
):
    activation_block = get_activation_block(activation)

    layers = [
        (
            "project",
            nn.Conv2d(input_channels, embedding_size, kernel_size=(1, 1), padding=0, bias=False),
        ),
    ]

    for layer_index in range(num_blocks):
        layers += [
            (
                f"conv{layer_index}",
                nn.Conv2d(embedding_size, embedding_size, kernel_size=(3, 3), padding=1, bias=False),
            ),
            (f"bn{layer_index}", nn.GroupNorm(num_groups, embedding_size)),
            (f"act{layer_index}", activation_block(inplace=True)),
        ]
    layers += [
        ("drop", nn.Dropout2d(dropout_rate, inplace=False)),
        ("conv1", nn.Conv2d(embedding_size, embedding_size, kernel_size=(1, 1), padding=0, bias=False)),
    ] + extra_final_layers

    return nn.Sequential(OrderedDict(layers))


class DecoupledHeadGroupNormLateShuffle(nn.Module):
    upsample_factor: int = 2

    def __init__(
        self,
        channels: int,
        classifier_dim: int,
        dropout_rate=0.25,
        objectness_dim: int = 32,
        size_dim: int = 32,
        offset_dim: int = 32,
        num_blocks: int = 3,
        activation: str = ACT_RELU,
        num_groups=16,
    ):
        super().__init__()
        self.objectness = build_regression_block(
            channels,
            embedding_size=objectness_dim,
            dropout_rate=dropout_rate,
            activation=activation,
            num_groups=num_groups,
            num_blocks=num_blocks,
            extra_final_layers=[
                (("shuffle"), nn.PixelShuffle(2)),
                ("final", nn.Conv2d(objectness_dim // 4, 1, kernel_size=(3, 3), padding=1)),
            ],
        )

        self.classifier = build_regression_block(
            channels,
            embedding_size=classifier_dim,
            dropout_rate=dropout_rate,
            num_groups=num_groups,
            activation=activation,
            num_blocks=num_blocks,
            extra_final_layers=[],
        )

        self.size = build_regression_block(
            channels,
            embedding_size=size_dim,
            dropout_rate=dropout_rate,
            activation=activation,
            num_groups=num_groups,
            num_blocks=num_blocks,
            extra_final_layers=[],
        )

        self.classifier_recombination = FeatureRecombination(classifier_dim, size_dim)
        self.classifier_up = nn.Sequential(
            OrderedDict(
                [
                    ("conv1", nn.Conv2d(classifier_dim, classifier_dim, kernel_size=(3, 3), padding=1)),
                    ("shuffle", nn.PixelShuffle(2)),
                    ("final", nn.Conv2d(classifier_dim // 4, 2, kernel_size=(3, 3), padding=1)),
                ]
            )
        )
        self.size_recombination = FeatureRecombination(size_dim, classifier_dim)
        self.size_up = nn.Sequential(
            OrderedDict(
                [
                    ("conv1", nn.Conv2d(size_dim, size_dim, kernel_size=(3, 3), padding=1)),
                    ("shuffle", nn.PixelShuffle(2)),
                    ("final", nn.Conv2d(size_dim // 4, 1, kernel_size=(3, 3), padding=1)),
                    ("activation", nn.Softplus()),
                ]
            )
        )

        if offset_dim:
            self.offset = build_regression_block(
                channels,
                embedding_size=offset_dim,
                dropout_rate=dropout_rate,
                activation=activation,
                num_groups=num_groups,
                num_blocks=num_blocks,
                extra_final_layers=[
                    (("shuffle"), nn.PixelShuffle(2)),
                    ("final", nn.Conv2d(offset_dim // 4, 2, kernel_size=(3, 3), padding=1)),
                ],
            )
        else:
            self.offset = None

        # Determine if offset is included
        if self.offset is not None:
            self.out_channels = 6  # Total output channels
        else:
            self.out_channels = 4  # Exclude offset channels

        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight)
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        # self.classifier_up.final.bias.data.fill_(0)
        # self.objectness.final.bias.data.fill_(-4)

        self.objectness.final.bias.data.fill_(-5)
        self.classifier_up.final.bias.data.fill_(-1)
        self.size_up.final.bias.data.fill_(2)

    def forward(self, features: Tensor) -> Dict[str, Tensor]:
        class_features = self.classifier(features)
        length_features = self.size(features)

        new_class_features = self.classifier_recombination(class_features, length_features)
        new_length_features = self.size_recombination(length_features, class_features)

        predicted_classes = self.classifier_up(new_class_features)
        predicted_length = self.size_up(new_length_features)

        output = {
            CENTERNET_OUTPUT_OBJECTNESS_MAP: self.objectness(features),
            CENTERNET_OUTPUT_VESSEL_MAP: predicted_classes[:, 0:1],
            CENTERNET_OUTPUT_FISHING_MAP: predicted_classes[:, 1:2],
            CENTERNET_OUTPUT_SIZE: predicted_length,
        }
        if self.offset is not None:
            output[CENTERNET_OUTPUT_OFFSET] = self.offset(features).clip(0, 1)
        return output