[DEBUG] Channel 1 shape: (2048, 2048)
(2007.0, 1786.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(2044.0, 1831.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(1856.0, 1345.0, 1.9375784397125244) [<class 'float'>, <class 'float'>, <class 'float'>]
Using encoder feature at index 3
x.shape after concat at stage 0: torch.Size([4, 688, 128, 128])
[DEBUG] Decoder stage 1
Using encoder feature at index 2
x.shape after concat at stage 1: torch.Size([4, 576, 256, 256])
[DEBUG] Decoder stage 2
Using encoder feature at index 1
x.shape after concat at stage 2: torch.Size([4, 424, 512, 512])
[DEBUG] Decoder stage 3
Using encoder feature at index 0
encoder_feature.shape after resizing at stage 3: torch.Size([4, 24, 1024, 1024])
x.shape after concat at stage 3: torch.Size([4, 280, 1024, 1024])
[DEBUG] feature_maps in FPNFuse: type=<class 'list'>
[DEBUG] Valid feature_maps[0] shape: torch.Size([4, 512, 128, 128])
[DEBUG] Valid feature_maps[1] shape: torch.Size([4, 384, 256, 256])
[DEBUG] Valid feature_maps[2] shape: torch.Size([4, 256, 512, 512])
[DEBUG] Valid feature_maps[3] shape: torch.Size([4, 128, 1024, 1024])
yolo_output shape: torch.Size([4, 21, 128, 128])
objectness_map shape: torch.Size([4, 1, 256, 256])
vessel_map shape: torch.Size([4, 1, 256, 256])
fishing_map shape: torch.Size([4, 1, 256, 256])
size_map shape: torch.Size([4, 1, 256, 256])
offset_map shape: torch.Size([4, 2, 256, 256])
yolo_output shape: torch.Size([4, 21, 128, 128])
classification_output shape: torch.Size([4, 2, 128, 128])
yolo_output upsampled to: torch.Size([4, 21, 256, 256])
classification_output upsampled to: torch.Size([4, 2, 256, 256])
Error executing job with overrides: ['+config=configs/multilabel_circle_net.yaml']
Traceback (most recent call last):
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 816, in <module>
    main()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 813, in main
    CenternetPipeline(cfg).train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 695, in train
    model = self.build_model(config, self.loaders)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 232, in build_model
    self.model = nas.search(
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 38, in search
    scores = [self.evaluation_function(model, train_loader, val_loader, device) for model in population]
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 38, in <listcomp>
    scores = [self.evaluation_function(model, train_loader, val_loader, device) for model in population]
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 665, in evaluate_architecture
    logger.debug(f"Outputs shape: {outputs.shape}, Encoded targets: {encoded_targets}")
AttributeError: 'dict' object has no attribute 'shape'
(sam3) samerkharboush@Samers-MBP SAR-MPS-OptunaNasFbnYoloTransformers % 

we have solved all the previous errors successfully

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, generate_yolo_targets
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
from torch.cuda.amp import autocast, GradScaler
from xview3.utils import TorchDevice
from xview3.evaluation import evaluate_on_scenes
from xview3.inference import predict_multilabel_scenes
from omegaconf import OmegaConf
from datetime import datetime

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.model: Optional[nn.Module] = None  # Ensure model is initialized
        self.criterion: Optional[nn.Module] = None 
        self.shore_root = os.path.join(os.environ.get('XVIEW3_DIR', ''), "full_dataset", "shoreline")
        self.device = TorchDevice.choose_torch_device()        
        self.feature_extractor = nn.Sequential(nn.Conv2d(3, 256, kernel_size=3, padding=1), nn.ReLU())
        self.yolo_head = YOLOHead(num_classes=2, anchors=[[10, 13], [16, 30], [33, 23]])  # Example anchors

    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
        )
        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']))}")
        self.model = model
        self.box_coder = MultilabelCircleNetCoder(
            image_size=(config.train.train_image_size[0], config.train.train_image_size[1]),
            output_stride=4,  
            max_objects=512,  
            heatmap_encoding="umich",
            labels_encoding="circle",
            ignore_value=-1
        )
        # Initialize ProgressiveNAS with the evaluation function
        nas = ProgressiveNAS(
            base_model=model,
            search_space=config.get('nas', {}).get('search_space', {}),
            evaluation_function=self.evaluate_architecture  # Ensure this method exists
        )

        # Search for the best model architecture
        self.model = nas.search(
            num_iterations=config.get('nas', {}).get('num_iterations', 4),
            train_loader=loaders['train'],
            val_loader=loaders['valid'],
            device=self.device
        )


        # Reinitialize YOLOHead and Classifier if needed
        self.model = YOLOHead(num_classes=self.cfg.model.num_classes, anchors=self.cfg.model.anchors).to(self.device)
        self.model.classifier = ConvolutionalClassifier(num_classes=config.model.num_classes)
        self.model.to(self.device)
        logger.debug(f"Model built: {self.model}")
        if self.model is None:
            self.model = YOLOHead(num_classes=self.cfg.model.num_classes, anchors=self.cfg.model.anchors).to(self.device)
            logger.debug("Model initialized successfully.")

        return self.model  # Return the best model found by NAS

    def build_criterions(self, config: Dict) -> Tuple[Dict[str, nn.Module], List[Callback]]:
        criterions, callbacks = super().build_criterions(config)
        
        self.criterion = YOLOLoss(num_classes=config.model.num_classes, anchors=config.model.anchors).to(self.device)
        if self.criterion is None:
            self.criterion = YOLOLoss(num_classes=config.model.num_classes, anchors=config.model.anchors).to(self.device)
            logger.debug("Loss criterion initialized successfully.")
        
        return {}, []

    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 if train_sampler else None,
            pin_memory=True,
            collate_fn=partial(centernet_collate, channels_last=train_ds.channels_last),
        )
        
        valid_loader = DataLoader(
            valid_ds,
            batch_size=config.get('batch_size', 8),
            num_workers=config.get('num_workers', 4),
            pin_memory=True,
            collate_fn=partial(centernet_collate, channels_last=valid_ds.channels_last),
        )

        return {"train": train_loader, "valid": valid_loader}, dataset_callbacks

    def calculate_bounding_boxes(self, predictions: pd.DataFrame) -> List[Dict[str, float]]:
        if predictions.empty:
            return []
        centers = predictions[['detect_scene_column', 'detect_scene_row']].values
        lengths = predictions['vessel_length_m'].values
        boxes = []
        for (cx, cy), length in zip(centers, lengths):
            if length > 0:
                half_length = length / 2
                box = {
                    'xmin': cx - half_length,
                    'ymin': cy - half_length,
                    'xmax': cx + half_length,
                    'ymax': cy + half_length,
                }
                boxes.append(box)
            else:
                boxes.append({'xmin': None, 'ymin': None, 'xmax': None, 'ymax': None})
        return boxes


    def encode_batch_targets(self, batch, device):
        scenes = batch.get('scenes')  # Get the scenes to run inference on
        if scenes is None or len(scenes) == 0:
            print("No scenes to process for predictions.")
            return {
                'boxes': torch.empty((0, 4), dtype=torch.float32, device=device),  # Return empty tensor for boxes
                'heatmap': torch.empty((0,), device=device),  # Adjust as per needed shape
                'is_vessel': torch.empty((0,), device=device),
                'is_fishing': torch.empty((0,), device=device),
                'lengths': torch.empty((0,), device=device),
                'offset': torch.empty((0,), device=device)
            }
        predictions = predict_multilabel_scenes(
            model=self.model,
            box_coder=self.box_coder,
            scenes=scenes,
            tile_step=32,
            tile_size=640,
            objectness_thresholds_lower_bound=0.5,
            normalization={},  # Specify normalization if required
            accumulate_on_gpu=False,
            fp16=False,
            batch_size=1,
            apply_activation=True,
            max_objects=100,
            channels_last=False,
            device=device,
            save_raw_predictions=True,  # This flag should be set as True/False as needed
            channels=['vv', 'vh']  # Pass an actual list of channel names here
        )
        
         # Ensure we handle empty predictions
        if predictions.empty:
            print("Predictions are empty; no vessels detected.")
            return {
                'boxes': torch.empty((0, 4), dtype=torch.float32, device=device),  # Return empty tensor for boxes
                'heatmap': torch.empty((0,), device=device),  # Adjust as per needed shape
                'is_vessel': torch.empty((0,), device=device),
                'is_fishing': torch.empty((0,), device=device),
                'lengths': torch.empty((0,), device=device),
                'offset': torch.empty((0,), device=device)
            }

        bounding_boxes = self.calculate_bounding_boxes(predictions)

        batch['boxes'] = torch.tensor(bounding_boxes, dtype=torch.float32, device=device)
        if 'boxes' in batch:
            gt_boxes = batch['boxes'].to(device)
        else:
            centers = batch.get('TARGET_CENTERS_KEY')  # Expected Shape: [batch_size, max_num_objs, 3]
            lengths = batch.get('TARGET_LENGTHS_KEY')    # Expected Shape: [batch_size, max_num_objs]
            is_vessel = batch.get('TARGET_VESSEL_KEY')   # Expected Shape: [batch_size, max_num_objs]
            is_fishing = batch.get('TARGET_FISHING_KEY') # Expected Shape: [batch_size, max_num_objs]

            # Ensure all necessary keys are present
            required_keys = ['TARGET_CENTERS_KEY', 'TARGET_LENGTHS_KEY', 'TARGET_VESSEL_KEY', 'TARGET_FISHING_KEY']
            if not all(k in batch for k in required_keys):
                missing = [k for k in required_keys if k not in batch]
                raise KeyError(f"Batch is missing required keys: {missing}")

            # Convert centers to numpy array and extract only [x, y]
            if isinstance(centers, list):
                # If centers are a list of tuples/lists
                centers = [c[:2] for c in centers]
            elif isinstance(centers, torch.Tensor):
                # If centers are a tensor, slice to get [x, y]
                centers = centers[:, :, :2].cpu().numpy()

            # Convert to NumPy array
            centers = np.array([np.array(c[:2]) for c in centers if isinstance(c, (list, tuple)) and len(c) >= 2])

            # Ensure lengths is of valid type
            if isinstance(lengths, torch.Tensor):
                lengths = lengths.cpu().numpy()
            elif isinstance(lengths, list):
                flattened_lengths = []
                for l in lengths:
                    # Check for empty tensors or lists
                    if isinstance(l, torch.Tensor) and l.numel() > 0:  # Only include non-empty tensors
                        flattened_lengths.extend(l.cpu().numpy())  # Prepare for flattening
                    elif isinstance(l, (list, np.ndarray)) and len(l) > 0:  # Include non-empty lists
                        flattened_lengths.extend(l)
            
                # Ensure flattened_lengths is a numpy array
                lengths = np.array(flattened_lengths, dtype=np.float32)

            # Ensure lengths is a 1D array
            if lengths.ndim != 1:
                raise ValueError(f"Expected lengths to be 1D, got {lengths.ndim}D.")

             # Convert centers to bounding boxes: [x_min, y_min, x_max, y_max]
            x_min = centers[..., 0] - (lengths / 2)
            y_min = centers[..., 1] - (lengths / 2)
            x_max = centers[..., 0] + (lengths / 2)
            y_max = centers[..., 1] + (lengths / 2)

            # Clamp coordinates to image boundaries
            img_width, img_height = self.box_coder.image_size  # Ensure this is defined appropriately
            x_min = np.clip(x_min, 0, img_width)
            y_min = np.clip(y_min, 0, img_height)
            x_max = np.clip(x_max, 0, img_width)
            y_max = np.clip(y_max, 0, img_height)

            # Stack to get boxes: [batch_size, max_num_objs, 4]
            boxes = np.stack([x_min, y_min, x_max, y_max], axis=-1)
            gt_boxes = torch.tensor(boxes, dtype=torch.float32, device=device)

            # Optionally, add 'boxes' to the batch dictionary for consistency
            batch['boxes'] = gt_boxes

        # Process other targets
        gt_confidences = batch.get('INPUT_CONFIDENCE_KEY', None)
        gt_classes = batch.get('TARGET_FISHING_KEY', None)  # Adjust based on your classification labels

        if gt_confidences is None or gt_classes is None:
            raise KeyError("Batch must contain 'INPUT_CONFIDENCE_KEY' and 'TARGET_FISHING_KEY'")

        # Ensure gt_confidences is a list of numerical values
        if isinstance(gt_confidences, torch.Tensor):
            gt_confidences = gt_confidences.cpu().numpy().tolist()
        elif isinstance(gt_confidences, np.ndarray):
            gt_confidences = gt_confidences.tolist()

        # Convert gt_classes to tensor if it's not
        if not isinstance(gt_classes, torch.Tensor):
            gt_classes = torch.tensor(gt_classes, dtype=torch.long, device=device)
        else:
            gt_classes = gt_classes.to(device)

        # Encode targets
        yolo_targets = self.box_coder.encode(
            centers=gt_boxes.cpu().numpy(),                # [batch_size, max_num_objs, 2] as NumPy array
            confidences=gt_confidences,                   # [batch_size, max_num_objs] as list of categorical labels
            lengths=lengths,                               # [batch_size, max_num_objs] as NumPy array
            is_vessel=is_vessel.cpu().numpy(),             # [batch_size, max_num_objs] as NumPy array
            is_fishing=is_fishing.cpu().numpy()            # [batch_size, max_num_objs] as NumPy array
        )

        # Convert encoded targets to tensors
        yolo_heatmap = torch.from_numpy(yolo_targets.heatmap).to(device)
        yolo_is_vessel = torch.from_numpy(yolo_targets.is_vessel).to(device)
        yolo_is_fishing = torch.from_numpy(yolo_targets.is_fishing).to(device)
        yolo_length = torch.from_numpy(yolo_targets.lengths).to(device)
        yolo_offset = torch.from_numpy(yolo_targets.offset).to(device)

        return {
            'heatmap': yolo_heatmap,
            'is_vessel': yolo_is_vessel,
            'is_fishing': yolo_is_fishing,
            'lengths': yolo_length,
            'offset': yolo_offset,
            'boxes': torch.tensor(bounding_boxes, dtype=torch.float32, device=device)  # Pass bounding boxes to your model
        }

    def train_step(self, batch, batch_idx):
        self.model.train()
        with autocast(device_type=self.device.type):
            images = batch['INPUT_IMAGE_KEY'].to(self.device)
            encoded_targets = self.encode_batch_targets(batch, self.device)
            outputs = self.model(images)
            loss = self.criterion(outputs, encoded_targets)

        # Scale the loss
        scaled_loss = loss / self.accumulation_steps
        self.scaler.scale(scaled_loss).backward()

        if (batch_idx + 1) % self.accumulation_steps == 0:
            self.scaler.step(self.optimizer)
            self.scaler.update()
            self.optimizer.zero_grad()

        return loss

    def evaluate_architecture(self, model, train_loader, val_loader, device):
        if self.model is None or self.criterion is None:
            raise RuntimeError("Model or criterion not initialized.")
        
        self.device = device
        self.model.eval()
        total_loss = 0.0
        total_samples = 0

        with torch.no_grad():
            for batch_idx, batch in enumerate(val_loader):
                images = batch['INPUT_IMAGE_KEY'].to(device)
                encoded_targets = self.encode_batch_targets(batch, device)

                if encoded_targets is None:
                    logger.warning("Encoded targets are None, skipping this batch.")
                    continue

                outputs = self.model(images)
                logger.debug(f"Outputs shape: {outputs.shape}, Encoded targets: {encoded_targets}")

                loss = self.criterion(outputs, encoded_targets)
                total_loss += loss.item() * images.size(0)
                total_samples += images.size(0)

        average_loss = total_loss / total_samples if total_samples > 0 else float('inf')  # Avoid division by zero
        print(f"[DEBUG] Average Evaluation Loss: {average_loss}")

        return -average_loss  


    def parse_evaluation_results(self, output_dir):
        """
        Implement parsing of evaluation results to extract metrics like average loss.
        This depends on how 'evaluate_on_scenes' outputs results.
        """
        # Example:
        thresholds_summary_path = os.path.join(output_dir, "nas_evaluation__step_4_thresholds_summary_.csv")
        if os.path.exists(thresholds_summary_path):
            thresholds_summary = pd.read_csv(thresholds_summary_path)
            average_loss = thresholds_summary['aggregate'].mean()
            return average_loss
        else:
            raise FileNotFoundError(f"Thresholds summary not found at {thresholds_summary_path}")

    def train(self):
        config = self.cfg
        self.loaders, dataset_callbacks = self.build_loaders(config)
        self.build_criterions(config)
        model = self.build_model(config, self.loaders)
        self.box_coder = model.box_coder 

        optimizer, scheduler, optimizer_callbacks = self.build_optimizer(model, self.loaders)
        # Initialize YOLOLoss with the correct number of classes and anchors
        num_classes = config.get('num_classes', 1)  # Adjust as per your dataset
        anchors = config.get('anchors', [[10,13], [16,30], [33,23]])  # Example anchors

        self.criterion = YOLOLoss(num_classes=num_classes, anchors=anchors).to(self.device)

        criterions, criterions_callbacks = self.build_criterions(config)
        metric_callbacks = self.build_metrics(config, self.loaders, model)

        nas = ProgressiveNAS(
            base_model=model,
            search_space=config.get('nas', {}).get('search_space', {}),
            evaluation_function=self.evaluate_architecture
        )
        self.model = nas.search(
            num_iterations=config.get('nas', {}).get('num_iterations', 4),
            train_loader=self.loaders['train'],
            val_loader=self.loaders['valid'],
            device=self.device
        )
        model = self.model  # Update reference to best_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)

        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 = TorchDevice.choose_torch_device()
        runner = SupervisedRunner(input_key=input_key, output_key=output_key, device=device)

        # Initialize GradScaler for mixed precision training
        self.scaler = GradScaler()

        # Set up gradient accumulation
        self.accumulation_steps = config.get('accumulation_steps', 1)

        extra_callbacks = [TimerCallback()]

        runner.train(
            fp16=self.build_distributed_params(config),
            model=model,
            criterion=self.criterion,  # Use the initialized YOLOLoss
            optimizer=optimizer,
            scheduler=scheduler,
            callbacks=dataset_callbacks + metric_callbacks + optimizer_callbacks + criterions_callbacks + extra_callbacks,
            loaders=self.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(),
            train_step=self.train_step,  # Ensure train_step is correctly referenced
        )

        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()

#test