DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x340939970>
[DEBUG] Channel 0 shape: (2048, 2048)
[DEBUG] Channel 1 shape: (2048, 2048)
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 735, 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 731, in main
    CenternetPipeline(cfg).train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 502, in train
    model = self.build_model(config, self.loaders)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 222, in build_model
    best_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 625, in evaluate_architecture
    centers = torch.tensor(centers, dtype=torch.float32)  # Choose appropriate dtype
ValueError: only one element tensors can be converted to Python scalars
(sam3) samerkharboush@Samers-MBP SAR-MPS-OptunaNasFbnYoloTransformers % 

here's some other files in my model might help you solve this issues
functional.py
import math
from typing import Tuple, Union

import numpy as np
import torch
from torch import Tensor

__all__ = [
    "bes_radius",
    "centernet_circle_2d",
    "centernet_gather_feat",
    "centernet_gaussian_2d",
    "centernet_heatmap_nms",
    "centernet_topk",
    "centernet_tranpose_and_gather_feat",
    "draw_umich_gaussian",
    "draw_msra_gaussian",
    "gaussian_radius",
    "pointwise_gaussian_2d",
    "centernet_tight_heatmap_nms",
    "compute_phi",
    "one_hot_labels",
    "change_box_order",
    "box_clamp",
    "box_iou",
    "box_select",
    "meshgrid",
    "bboxes_area_xyxy",
]


def bboxes_area_xyxy(bboxes: Union[Tensor, np.ndarray]) -> Union[Tensor, np.ndarray]:
    """

    Args:
        bboxes: Tensor of bboxes [N,4] in xyxy format

    Returns:

    """
    w = bboxes[:, 2] - bboxes[:, 0]
    h = bboxes[:, 3] - bboxes[:, 1]
    return w * h


def centernet_gaussian_2d(shape, sigma=1.0):
    m, n = [(ss - 1.0) / 2.0 for ss in shape]
    y, x = np.ogrid[-m : m + 1, -n : n + 1]

    h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
    h[h < np.finfo(h.dtype).eps * h.max()] = 0
    return h


def centernet_circle_2d(r: int, value: float, dtype):
    a, b = r, r
    n = 2 * r + 1

    y, x = np.ogrid[-a : n - a, -b : n - b]
    mask = x * x + y * y <= r * r

    array = np.zeros((n, n), dtype=dtype)
    value = np.array(value, dtype=dtype).reshape((1, 1))
    array[:] = value
    array[~mask] = 0
    return array


def pointwise_gaussian_2d():
    pos_kernel = np.float32([[0.5, 0.75, 0.5], [0.75, 1.0, 0.75], [0.5, 0.75, 0.5]])
    return pos_kernel


def gaussian_radius(det_size, min_overlap=0.7):
    height, width = det_size

    a1 = 1
    b1 = height + width
    c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
    sq1 = np.sqrt(b1 ** 2 - 4 * a1 * c1)
    r1 = (b1 + sq1) / 2

    a2 = 4
    b2 = 2 * (height + width)
    c2 = (1 - min_overlap) * width * height
    sq2 = np.sqrt(b2 ** 2 - 4 * a2 * c2)
    r2 = (b2 + sq2) / 2

    a3 = 4 * min_overlap
    b3 = -2 * min_overlap * (height + width)
    c3 = (min_overlap - 1) * width * height
    sq3 = np.sqrt(b3 ** 2 - 4 * a3 * c3)
    r3 = (b3 + sq3) / 2
    r = min(r1, r2, r3)
    return max(1, math.ceil(r))


def compute_phi(w, h):
    a, b = min(w, h), max(w, h)
    aspect_ratio = a / b
    angle = 45 * aspect_ratio
    if w > h:
        angle = 90 - angle
    return angle


def bes_radius(det_size, min_overlap=0.5) -> int:
    """Compute radius of gaussian.
    Arguments:
        w (int): weight of box.
        h (int): height of box.
        iou_threshold (float): min required IOU between gt and smoothed box.
    Returns:
        radius (int): radius of gaussian.
    """
    w, h = det_size
    phi = compute_phi(w, h)
    sin_phi = math.sin(math.radians(phi))
    cos_phi = math.cos(math.radians(phi))
    a = sin_phi * cos_phi
    b = -(w * sin_phi + h * cos_phi)
    c = w * h * (1 - min_overlap) / (1 + min_overlap)
    d = math.sqrt(b * b - 4 * a * c)
    r = -(b + d) / (2 * a)
    return int(max(1, math.ceil(r)))


def gaussian2DNonSquare(size_x, size_y, alpha, dtype=np.float32):
    y, x = np.ogrid[0:size_y, 0:size_x]
    mx = size_x // 2
    my = size_y // 2
    sigma_x = alpha * size_x / 6
    sigma_y = alpha * size_y / 6
    h = np.exp(-((x - mx) ** 2.0 / (2.0 * sigma_x ** 2.0) + (y - my) ** 2.0 / (2.0 * sigma_y ** 2.0))).astype(dtype)
    h[h < np.finfo(h.dtype).eps * h.max()] = 0
    area = np.sum(h > 0)
    return h, area


def draw_msra_gaussian(heatmap, center, sigma):
    tmp_size = sigma * 3
    mu_x = int(center[0] + 0.5)
    mu_y = int(center[1] + 0.5)
    w, h = heatmap.shape[0], heatmap.shape[1]
    ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]
    br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]
    if ul[0] >= h or ul[1] >= w or br[0] < 0 or br[1] < 0:
        return heatmap
    size = 2 * tmp_size + 1
    x = np.arange(0, size, 1, np.float32)
    y = x[:, np.newaxis]
    x0 = y0 = size // 2
    g = np.exp(-((x - x0) ** 2 + (y - y0) ** 2) / (sigma ** 2))
    g_x = max(0, -ul[0]), min(br[0], h) - ul[0]
    g_y = max(0, -ul[1]), min(br[1], w) - ul[1]
    img_x = max(0, ul[0]), min(br[0], h)
    img_y = max(0, ul[1]), min(br[1], w)
    heatmap[img_y[0] : img_y[1], img_x[0] : img_x[1]] = np.maximum(
        heatmap[img_y[0] : img_y[1], img_x[0] : img_x[1]], g[g_y[0] : g_y[1], g_x[0] : g_x[1]]
    )
    return heatmap


def draw_umich_gaussian(heatmap: np.ndarray, center: Tuple[int, int], radius: int) -> np.ndarray:
    if radius == "pointwise":
        gaussian = pointwise_gaussian_2d()
        radius = 1
    else:
        diameter = 2 * radius + 1
        gaussian = centernet_gaussian_2d((diameter, diameter), sigma=diameter / 6.0)

    x, y = int(center[0]), int(center[1])

    height, width = heatmap.shape[0:2]

    left, right = min(x, radius), min(width - x, radius + 1)
    top, bottom = min(y, radius), min(height - y, radius + 1)

    masked_heatmap = heatmap[y - top : y + bottom, x - left : x + right]
    masked_gaussian = gaussian[radius - top : radius + bottom, radius - left : radius + right]
    if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0:
        np.maximum(masked_heatmap, masked_gaussian, out=masked_heatmap)
    return heatmap


def draw_circle(size_map: np.ndarray, center: Tuple[int, int], radius: int, value: float):
    """
    Draw bbox size or offset in a radius
    :param size_map: [2, rows, cols]
    :param center: (x,y)
    :param radius: int
    :param value: [2]
    :return:
    """
    gaussian = centernet_circle_2d(radius, value, dtype=size_map.dtype)

    x, y = int(center[0]), int(center[1])

    height, width = size_map.shape

    left, right = min(x, radius), min(width - x, radius + 1)
    top, bottom = min(y, radius), min(height - y, radius + 1)

    masked_heatmap = size_map[y - top : y + bottom, x - left : x + right]
    masked_gaussian = gaussian[radius - top : radius + bottom, radius - left : radius + right]
    if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0:
        np.maximum(masked_heatmap, masked_gaussian, out=masked_heatmap)
    return size_map


def disk3_max_pool2d(heatmap):
    hp = torch.max_pool2d(heatmap, kernel_size=(3, 1), padding=(1, 0), stride=1)
    vp = torch.max_pool2d(heatmap, kernel_size=(1, 3), padding=(0, 1), stride=1)
    keep = (heatmap == hp) & (heatmap == vp)
    return heatmap * keep.type_as(heatmap)


def disk5_max_pool2d(heatmap):
    hp = torch.max_pool2d(heatmap, kernel_size=(5, 1), padding=(2, 0), stride=1)
    vp = torch.max_pool2d(heatmap, kernel_size=(1, 5), padding=(0, 2), stride=1)
    mp = torch.max_pool2d(heatmap, kernel_size=(3, 3), padding=(1, 1), stride=1)
    keep = (heatmap == hp) & (heatmap == vp) & (heatmap == mp)
    return heatmap * keep.type_as(heatmap)


def centernet_heatmap_nms(heatmap: Tensor, kernel=3) -> Tensor:
    pad = (kernel - 1) // 2

    hmax = torch.max_pool2d(heatmap, (kernel, kernel), stride=1, padding=pad)
    keep = (hmax == heatmap).type_as(heatmap)
    return heatmap * keep


def centernet_tight_heatmap_nms(heatmap: Tensor, kernel=3) -> Tensor:
    if kernel == 3:
        return disk3_max_pool2d(heatmap)
    elif kernel == 5:
        return disk5_max_pool2d(heatmap)
    else:
        raise KeyError(f"Unsupported kernel size {kernel}")


def centernet_topk(scores: Tensor, top_k: int = 100) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
    batch, cat, height, width = scores.size()

    topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), top_k)

    topk_inds = topk_inds % (height * width)
    topk_ys = torch.div(topk_inds, width, rounding_mode="trunc").int().float()
    topk_xs = (topk_inds % width).int().float()

    topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), top_k)
    topk_clses = torch.div(topk_ind, top_k, rounding_mode="trunc").int()
    topk_inds = centernet_gather_feat(topk_inds.view(batch, -1, 1), topk_ind).view(batch, top_k)
    topk_ys = centernet_gather_feat(topk_ys.view(batch, -1, 1), topk_ind).view(batch, top_k)
    topk_xs = centernet_gather_feat(topk_xs.view(batch, -1, 1), topk_ind).view(batch, top_k)

    return topk_score, topk_inds, topk_clses, topk_ys, topk_xs


def centernet_gather_feat(feat, ind, mask=None):
    dim = feat.size(2)
    ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)
    feat = feat.gather(1, ind)
    if mask is not None:
        mask = mask.unsqueeze(2).expand_as(feat)
        feat = feat[mask]
        feat = feat.view(-1, dim)
    return feat


def centernet_tranpose_and_gather_feat(feat, ind):
    feat = feat.permute(0, 2, 3, 1).contiguous()
    feat = feat.view(feat.size(0), -1, feat.size(3))
    feat = centernet_gather_feat(feat, ind)
    return feat


def one_hot_labels(tensor, num_classes: int):
    """Encodes class labes as one-hot vector. Negative labels converted too background class.

    :param tensor: Labels vector of shape [N]
    :param num_classes: Number of classes, including dummy background class. E.g num_real_classes + 1
    :return: Labels tensor of shape [N, num_classes]
    """
    if len(tensor) == 0:
        return torch.from_numpy(np.array([], dtype=int)).long()

    tensor = tensor.clone()
    tensor[tensor < 0] = 0
    one_hot_tensor = torch.zeros(tensor.size() + torch.Size([num_classes]))
    dim = len(tensor.size())
    index = tensor.unsqueeze(dim)
    return one_hot_tensor.scatter_(dim, index, 1)


def change_box_order(boxes, order, dim=-1):
    """Change box order between (x_min, y_min, x_max, y_max) and (x_center, y_center, width, height).

    Args:
      boxes: (tensor) bounding boxes, sized [N, 4].
      order: (str) either 'xyxy2xywh' or 'xywh2xyxy'.

    Returns:
      (tensor) converted bounding boxes, sized [N, 4].
    """
    assert order in ["xyxy2xywh", "xywh2xyxy"]
    a = boxes[..., 0:2]
    b = boxes[..., 2:4]
    if order == "xyxy2xywh":
        return torch.cat([(a + b) / 2, b - a], dim)
    return torch.cat([a - b / 2, a + b / 2], dim)


def box_clamp(boxes, x_min, y_min, x_max, y_max):
    """Clamp boxes.

    Args:
      boxes: (tensor) bounding boxes of (x_min, y_min, x_max, y_max), sized [N, 4].
      x_min: (number) min value of x.
      y_min: (number) min value of y.
      x_max: (number) max value of x.
      y_max: (number) max value of y.

    Returns:
      (tensor) clamped boxes.
    """
    boxes[:, 0] = boxes[:, 0].clamp(min=x_min, max=x_max)
    boxes[:, 1] = boxes[:, 1].clamp(min=y_min, max=y_max)
    boxes[:, 2] = boxes[:, 2].clamp(min=x_min, max=x_max)
    boxes[:, 3] = boxes[:, 3].clamp(min=y_min, max=y_max)
    return boxes


def box_select(boxes, x_min, y_min, x_max, y_max):
    """Select boxes in range (x_min, y_min, x_max, y_max).

    Args:
      boxes: (tensor) bounding boxes of (x_min, y_min, x_max, y_max), sized [N, 4].
      x_min: (number) min value of x.
      y_min: (number) min value of y.
      x_max: (number) max value of x.
      y_max: (number) max value of y.

    Returns:
      (tensor) selected boxes, sized [M, 4].
      (tensor) selected mask, sized [N, ].
    """
    mask = (boxes[:, 0] >= x_min) & (boxes[:, 1] >= y_min) & (boxes[:, 2] <= x_max) & (boxes[:, 3] <= y_max)
    boxes = boxes[mask, :]
    return boxes, mask


def box_iou(box1, box2):
    """Compute the intersection over union of two set of boxes.

    The box order must be (x_min, y_min, x_max, y_max).

    Args:
      box1: (tensor) bounding boxes, sized [N, 4].
      box2: (tensor) bounding boxes, sized [M, 4].

    Return:
      (tensor) iou, sized [N, M].

    Reference:
      https://github.com/chainer/chainercv/blob/master/chainercv/utils/bbox/bbox_iou.py
    """
    # N = box1.size(0)
    # M = box2.size(0)

    lt = torch.max(box1[:, None, :2], box2[:, :2])  # [N, M, 2]
    rb = torch.min(box1[:, None, 2:], box2[:, 2:])  # [N, M, 2]

    wh = (rb - lt).clamp(min=0)  # [N, M, 2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N, M]

    area1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])  # [N, ]
    area2 = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1])  # [M, ]
    iou = inter / (area1[:, None] + area2 - inter)
    return iou


def meshgrid(x, y, row_major=True):
    """Return meshgrid in range x & y.

    Args:
      x: (int) first dim range.
      y: (int) second dim range.
      row_major: (bool) row major or column major.

    Returns:
      (tensor) meshgrid, sized [x * y, 2]

    Example:
    >> meshgrid(3, 2)
    0  0
    1  0
    2  0
    0  1
    1  1
    2  1
    [torch.FloatTensor of size 6x2]

    >> meshgrid(3, 2, row_major=False)
    0  0
    0  1
    0  2
    1  0
    1  1
    1  2
    [torch.FloatTensor of size 6x2]
    """
    a = torch.arange(0, x, dtype=torch.float)
    b = torch.arange(0, y, dtype=torch.float)
    xx = a.repeat(y).view(-1, 1)
    yy = b.view(-1, 1).repeat(1, x).view(-1, 1)
    return torch.cat([xx, yy], 1) if row_major else torch.cat([yy, xx], 1)


#boxer_dropout.py
import random

import albumentations as A
import cv2
import numpy as np
import torch

__all__ = ["BoxesDropout", "CoarseDropoutWithBboxes"]


class BoxesDropout(A.DualTransform):
    """
    Remove objects and fill image & mask corresponding to removed bboxes.
    """

    def __init__(
        self,
        max_objects=5,
        max_fraction=0.2,
        image_fill_value=0,
        mask_fill_value=0,
        always_apply=False,
        drop_overlapping_boxes=True,
        overlap_iou=0.35,
        p=0.5,
    ):
        """
        Args:
            max_objects: Maximum number of labels that can be zeroed out. Can be tuple, in this case it's [min, max]
            image_fill_value: Fill value to use when filling image.
                Can be 'inpaint' to apply inpaining (works only  for 3-chahnel images)
            mask_fill_value: Fill value to use when filling mask.

        Targets:
            image, mask

        Image types:
            uint8, float32
        """
        super(BoxesDropout, self).__init__(always_apply, p)
        self.max_objects = max_objects
        self.max_fraction = max_fraction
        self.image_fill_value = image_fill_value
        self.mask_fill_value = mask_fill_value
        self.drop_overlapping_boxes = drop_overlapping_boxes
        self.overlap_iou = overlap_iou

    @property
    def targets_as_params(self):
        return ["image", "bboxes"]

    @property
    def targets(self):
        return {
            "image": self.apply,
            "mask": self.apply_to_mask,
            "masks": self.apply_to_masks,
            "bboxes": self.apply_to_bboxes,
            "keypoints": self.apply_to_keypoints,
        }

    def get_params_dependent_on_targets(self, params):
        from torchvision.ops import box_iou

        image = params["image"]
        rows, cols = image.shape[:2]
        bboxes = A.denormalize_bboxes(params["bboxes"], rows, cols)

        num_bboxes = len(bboxes)
        max_num_objects_to_drop = min(self.max_objects, int(self.max_fraction * num_bboxes))

        if max_num_objects_to_drop == 0:
            dropout_mask = None
            objects_to_drop = []
        else:
            indexes = np.arange(num_bboxes)

            objects_to_drop = random.randint(1, max_num_objects_to_drop)
            objects_to_drop = set(random.sample(indexes.tolist(), objects_to_drop))

            if self.drop_overlapping_boxes:
                bboxes = np.asarray([box[:4] for box in bboxes]).reshape((-1, 4))
                iou: np.ndarray = box_iou(torch.from_numpy(bboxes), torch.from_numpy(bboxes)).numpy()
                np.fill_diagonal(iou, 0)

                additional_objects_to_drop = []

                for obj in objects_to_drop:
                    overlapping = np.flatnonzero(iou[obj] > self.overlap_iou)
                    additional_objects_to_drop.extend(overlapping)

                objects_to_drop.union(set(additional_objects_to_drop))

                if len(objects_to_drop) > max_num_objects_to_drop:
                    # If the total number of objects to drop exceeds the allowed number,
                    # we revert skip dropout
                    objects_to_drop = []

            dropout_mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.bool)
            for object_index, object_bbox in enumerate(bboxes):
                x_min, y_min, x_max, y_max = object_bbox[:4]
                if object_index in objects_to_drop:
                    dropout_mask[int(y_min) : int(y_max), int(x_min) : int(x_max)] = True

        return {"dropout_mask": dropout_mask, "objects_to_drop": objects_to_drop}

    def apply_to_bboxes(self, bboxes, **params):
        objects_to_drop = params["objects_to_drop"]
        return [box for i, box in enumerate(bboxes) if i not in objects_to_drop]

    def apply(self, img, dropout_mask=None, **params):
        if dropout_mask is None:
            return img

        if self.image_fill_value == "inpaint":
            dropout_mask = dropout_mask.astype(np.uint8)
            _, _, w, h = cv2.boundingRect(dropout_mask)
            radius = min(3, max(w, h) // 2)
            img = cv2.inpaint(img, dropout_mask, radius, cv2.INPAINT_NS)
        else:
            img = img.copy()
            img[dropout_mask] = self.image_fill_value

        return img

    def apply_to_mask(self, img, dropout_mask=None, **params):
        if dropout_mask is None:
            return img

        img = img.copy()
        img[dropout_mask] = self.mask_fill_value
        return img

    def get_transform_init_args_names(self):
        return ("max_objects", "image_fill_value", "mask_fill_value")


class CoarseDropoutWithBboxes(A.CoarseDropout):
    def __init__(
        self,
        uncertain_visibility: float = 0.75,
        uncertain_label: int = None,
        min_visibility=0.5,
        max_holes=8,
        max_height=8,
        max_width=8,
        min_holes=None,
        min_height=None,
        min_width=None,
        fill_value=0,
        mask_fill_value=None,
        always_apply=False,
        p=0.5,
    ):
        if uncertain_visibility and uncertain_label is None:
            raise ValueError()

        super().__init__(
            max_holes=max_holes,
            max_height=max_height,
            max_width=max_width,
            min_holes=min_holes,
            min_height=min_height,
            min_width=min_width,
            fill_value=fill_value,
            mask_fill_value=mask_fill_value,
            always_apply=always_apply,
            p=p,
        )
        self.uncertain_label = uncertain_label
        self.uncertain_visibility = uncertain_visibility
        self.min_visibility = min_visibility

    @property
    def targets_as_params(self):
        return ["image", "bboxes"]

    def get_params_dependent_on_targets(self, params):
        img = params["image"]
        bboxes = params["bboxes"]

        rows, cols = img.shape[:2]

        holes = []

        for _n in range(random.randint(self.min_holes, self.max_holes)):
            hole_height = random.randint(self.min_height, self.max_height)
            hole_width = random.randint(self.min_width, self.max_width)

            y1 = random.randint(0, rows - hole_height)
            x1 = random.randint(0, cols - hole_width)
            y2 = y1 + hole_height
            x2 = x1 + hole_width
            holes.append((x1, y1, x2, y2))

        holes_mask = np.zeros((rows, cols), dtype=np.bool)
        for x1, y1, x2, y2 in holes:
            holes_mask[y1:y2, x1:x2] = True

        keep_bboxes = []
        new_labels = []

        for i, bbox in enumerate(bboxes):
            x1, y1, x2, y2, label = A.denormalize_bbox(bbox, rows, cols)
            x1 = int(x1)
            y1 = int(y1)
            x2 = int(x2)
            y2 = int(y2)

            bbox_mask = np.zeros((rows, cols), dtype=np.bool)
            bbox_mask[y1:y2, x1:x2] = True
            area_before = bbox_mask.sum()
            bbox_mask[holes_mask] = False
            area_after = bbox_mask.sum()
            visibility = area_after / area_before

            if visibility > self.min_visibility:
                if self.uncertain_visibility is not None and visibility < self.uncertain_visibility:
                    label = self.uncertain_label

            keep_bboxes.append(visibility > self.min_visibility)
            new_labels.append(label)

        return {"holes": holes, "keep_bboxes": keep_bboxes, "new_labels": new_labels}

    def apply_to_bboxes(self, bboxes, keep_bboxes=None, new_labels=None, **params):
        bboxes = [tuple(box[:4]) + tuple([label]) for box, label, keep in zip(bboxes, new_labels, keep_bboxes) if keep]
        return bboxes

#coarse_dropout.py
from typing import Tuple

import albumentations as A

__all__ = ["CoarseDropout"]


class CoarseDropout(A.CoarseDropout):
    def _keypoint_in_hole(self, keypoint: Tuple, hole: Tuple) -> bool:
        x1, y1, x2, y2 = hole
        x, y = keypoint[:2]
        return x1 <= x < x2 and y1 <= y < y2

    def apply_to_keypoints(self, keypoints, holes=(), **params):
        for hole in holes:
            remaining_keypoints = []
            for kp in keypoints:
                if not self._keypoint_in_hole(kp, hole):
                    remaining_keypoints.append(kp)
            keypoints = remaining_keypoints
        return keypoints


def main():
    from matplotlib import pyplot as plt
    import cv2
    import numpy as np

    t = A.Compose(
        [CoarseDropout(p=1, min_width=32, min_height=32, max_width=128, max_height=128)],
        keypoint_params=A.KeypointParams(format="xy", label_fields=["labels"]),
    )

    # image = np.zeros((256, 256, 3), dtype=np.uint8)
    image = cv2.imread("lena.png")[..., ::-1].copy()
    # image[::10, :] = 255
    # image[::35, :] = 255
    # image[::51, :] = 255
    # image[:, ::11] = 255
    # image[:, ::24] = 255
    # image[:, ::45] = 255

    kpts = []
    for x in range(20, 512, 20):
        for y in range(20, 512, 20):
            kpts.append((x, y))

    for k in kpts:
        image = cv2.circle(image, k, radius=3, color=(0, 200, 0), thickness=2, lineType=cv2.LINE_AA)

    data = t(image=image, keypoints=kpts, labels=np.zeros(len(kpts)))
    print(kpts)
    print(data["keypoints"])

    for k in data["keypoints"]:
        data["image"] = cv2.circle(
            data["image"],
            tuple(map(int, k[:2])),
            radius=2,
            color=(200, 0, 200),
            thickness=2,
            lineType=cv2.LINE_AA,
        )

    f, ax = plt.subplots(2, 1, figsize=(6, 12))
    ax[0].imshow(image)
    ax[1].imshow(data["image"])
    f.tight_layout()
    f.show()


if __name__ == "__main__":
    main()

#copy_paste.py
import math
import random
from typing import List, Union, Tuple, Callable

import albumentations as A
import cv2
import numpy as np
from sklearn.utils import compute_sample_weight

__all__ = ["CopyPasteAugmentation"]


class CopyPasteAugmentation(A.DualTransform):
    def __init__(
        self,
        images: List[str],
        bboxes: List,
        labels: List[int],
        transform: A.Compose,
        read_image_fn: Callable,
        class_weights: Union[str, np.ndarray] = "balanced",
        seamless_clone_p=0,
        always_apply=False,
        p=0.5,
    ):
        super().__init__(always_apply=always_apply, p=p)
        self.seamless_clone_p = seamless_clone_p
        self.images = images
        self.bboxes = np.asarray(bboxes, dtype=int)
        self.labels = np.asarray(labels, dtype=int)
        self.transform = transform
        self.read_image = read_image_fn
        if class_weights == "balanced":
            sample_weights = compute_sample_weight("balanced", labels)
        else:
            sample_weights = compute_sample_weight(class_weights, labels)

        self.sample_weights = sample_weights

    @property
    def targets(self):
        return {
            "image": self.apply,
            "bboxes": self.apply_to_bboxes,
        }

    @property
    def targets_as_params(self):
        return "image", "bboxes"

    def get_params_dependent_on_targets(self, params):
        image = params["image"]
        rows, cols = image.shape[:2]
        bboxes = params["bboxes"]
        bboxes = A.convert_bboxes_from_albumentations(bboxes, "pascal_voc", rows, cols)

        # Compute average object size
        if len(bboxes) != 0:
            bboxes = np.array(bboxes)[:, :4]
            median_size = np.median(np.sqrt((bboxes[:, 2] - bboxes[:, 0]) * (bboxes[:, 3] - bboxes[:, 1])))
        else:
            median_size = None

        seed_image, seed_label = self._select_box()

        if median_size is not None:
            seed_size = math.sqrt(seed_image.shape[0] * seed_image.shape[1])
            scale = min(6.0, max(0.1, random.gauss(median_size / seed_size, 0.5)))
            seed_image = cv2.resize(seed_image, dsize=None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)

        # Paint used regions
        mask = np.ones((rows, cols), dtype=np.uint8) * 255
        for (x1, y1, x2, y2) in bboxes:
            mask[int(y1) : int(y2), int(x1) : int(x2)] = 0

        mask = cv2.distanceTransform(mask, cv2.DIST_L2, 3, dstType=cv2.CV_32F)

        max_size = max(seed_image.shape[0], seed_image.shape[1])
        half_width = seed_image.shape[1] // 2
        half_height = seed_image.shape[0] // 2

        mask[: half_height + 1, :] = 0
        mask[:, : half_width + 1] = 0
        mask[mask.shape[0] - half_height - 1 :, :] = 0
        mask[:, mask.shape[1] - half_width - 1 :] = 0

        local_max = mask > max_size
        if not local_max.any():
            return {}

        nz_rows, nz_cols = np.nonzero(local_max)

        index = random.choice(np.arange(len(nz_rows)))

        x1 = nz_cols[index] - seed_image.shape[1] // 2
        y1 = nz_rows[index] - seed_image.shape[0] // 2
        x2 = x1 + seed_image.shape[1]
        y2 = y1 + seed_image.shape[0]

        return {
            "seed_image": seed_image,
            "seed_bbox": (x1, y1, x2, y2),
            "seed_p": ((x1 + x2) // 2, (y1 + y2) // 2),
            "seed_label": seed_label,
            "use_seamless_clone": self.seamless_clone_p > random.random(),
        }

    def apply(self, img, seed_image=None, seed_p=None, use_seamless_clone=False, seed_bbox=None, **params):
        if seed_image is not None:
            if use_seamless_clone:
                mask = np.ones(seed_image.shape[:2], dtype=np.uint8) * 255
                mask[0, :] = 0
                mask[:, 0] = 0
                mask[mask.shape[0] - 1, :] = 0
                mask[:, mask.shape[1] - 1] = 0

                return cv2.seamlessClone(
                    src=np.ascontiguousarray(seed_image),
                    dst=np.ascontiguousarray(img),
                    mask=mask,
                    p=seed_p,
                    flags=cv2.NORMAL_CLONE,
                )
            else:
                x1, y1, x2, y2 = seed_bbox
                img_hard = img.copy()
                img_hard[y1:y2, x1:x2] = seed_image
                return img_hard
        return img

    def apply_to_bboxes(self, bboxes, seed_bbox=None, seed_label=None, **params):
        if seed_bbox is not None:
            t = A.convert_bbox_to_albumentations(seed_bbox, "pascal_voc", params["rows"], params["cols"])
            bboxes = bboxes + [(*t, seed_label)]
        return bboxes

    def _select_box(self) -> Tuple[np.ndarray, int]:
        n = len(self.images)
        (index,) = random.choices(np.arange(n), self.sample_weights, k=1)

        image = self.read_image(self.images[index])
        x1, y1, x2, y2 = self.bboxes[index]

        roi = slice(y1, y2), slice(x1, x2)
        seed_image = image[roi]
        seed_label = self.labels[index]

        # Augment
        seed_image = self.transform(image=seed_image)["image"]

        return seed_image, seed_label

    def get_transform_init_args_names(self):
        return ()


#elastic_transform.py
import math
import random

import albumentations as A
import cv2
import matplotlib.pyplot as plt
import numpy as np
#from albumentations.augmentations.functional import maybe_process_in_chunks as _maybe_process_in_chunks
from albumentations.augmentations.functional import _maybe_process_in_chunks
__all__ = ["ElasticTransform"]


@A.preserve_channel_dim
def elastic_transform(
    img,
    map_x,
    map_y,
    interpolation=cv2.INTER_LINEAR,
    border_mode=cv2.BORDER_REFLECT_101,
    value=None,
):
    remap_fn = _maybe_process_in_chunks(
        cv2.remap,
        map1=map_x,
        map2=map_y,
        interpolation=interpolation,
        borderMode=border_mode,
        borderValue=value,
    )
    return remap_fn(img)


class ElasticTransform(A.DualTransform):
    def __init__(
        self,
        alpha=1,
        sigma=50,
        interpolation=cv2.INTER_LINEAR,
        border_mode=cv2.BORDER_REFLECT_101,
        value=None,
        mask_value=None,
        always_apply=False,
        p=0.5,
    ):
        super().__init__(always_apply=always_apply, p=p)
        self.alpha = A.to_tuple(alpha)
        self.sigma = A.to_tuple(sigma)
        self.interpolation = interpolation
        self.border_mode = border_mode
        self.value = value
        self.mask_value = mask_value

    def get_transform_init_args_names(self):
        return ("alpha", "sigma", "interpolation", "border_mode", "value", "mask_value")

    def apply(self, img, sigma=0, alpha=0, map_x=None, map_y=None, interpolation=cv2.INTER_LINEAR, **params):
        return elastic_transform(
            img,
            map_x=map_x,
            map_y=map_y,
            interpolation=interpolation,
            border_mode=self.border_mode,
            value=self.value,
        )

    def apply_to_mask(self, img, sigma=0, alpha=0, map_x=None, map_y=None, **params):
        return elastic_transform(
            img,
            map_x=map_x,
            map_y=map_y,
            interpolation=cv2.INTER_NEAREST,
            border_mode=self.border_mode,
            value=self.value,
        )

    def update_params(self, params, **kwargs):
        height, width = kwargs["image"].shape[:2]

        dx = np.zeros((height, width))
        dy = np.zeros((height, width))

        for _ in range(128):
            dx[random.randrange(0, height), random.randrange(0, width)] = random.uniform(self.alpha[0], self.alpha[1])
            dy[random.randrange(0, height), random.randrange(0, width)] = random.uniform(self.alpha[0], self.alpha[1])

        for _ in range(32):
            dx = cv2.blur(dx, (7, 7))
            dy = cv2.blur(dy, (7, 7))

        x, y = np.meshgrid(np.arange(width), np.arange(height))

        params["map_x"] = np.float32(x + dx)
        params["map_y"] = np.float32(y + dy)
        return params

    def apply_to_keypoint(self, keypoint, **params):
        x, y = keypoint[:2]
        map_x, map_y = params["map_x"], params["map_y"]
        mask = np.zeros(map_x.shape[:2], dtype=np.uint8)
        mask[y, x] = 255
        mask = cv2.remap(mask, map_x, map_y, borderMode=cv2.BORDER_CONSTANT, borderValue=0, interpolation=cv2.INTER_LINEAR)
        _, _, _, maxLoc = cv2.minMaxLoc(mask)
        xn, yn = maxLoc
        return (xn, yn) + keypoint[2:]


def main():
    t = A.Compose(
        [ElasticTransform(p=1, alpha=(1000, 5000))],
        keypoint_params=A.KeypointParams(format="xy", label_fields=["labels"]),
    )

    # image = np.zeros((256, 256, 3), dtype=np.uint8)
    image = cv2.imread("lena.png")[..., ::-1].copy()
    # image[::10, :] = 255
    # image[::35, :] = 255
    # image[::51, :] = 255
    # image[:, ::11] = 255
    # image[:, ::24] = 255
    # image[:, ::45] = 255

    kpts = []
    for x in range(20, 512, 20):
        for y in range(20, 512, 20):
            kpts.append((x, y))

    for k in kpts:
        image = cv2.circle(image, k, radius=3, color=(0, 200, 0), thickness=2, lineType=cv2.LINE_AA)

    data = t(image=image, keypoints=kpts, labels=np.zeros(len(kpts)))
    print(kpts)
    print(data["keypoints"])

    for k in data["keypoints"]:
        data["image"] = cv2.circle(
            data["image"],
            tuple(map(int, k[:2])),
            radius=2,
            color=(200, 0, 200),
            thickness=2,
            lineType=cv2.LINE_AA,
        )

    f, ax = plt.subplots(2, 1, figsize=(6, 12))
    ax[0].imshow(image)
    ax[1].imshow(data["image"])
    f.tight_layout()
    f.show()


if __name__ == "__main__":
    main()

#and the upgraded train_multilabel.py which has the error mismatch
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 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.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']))}")

        # 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
        best_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
        model.yolo_head = YOLOHead(num_classes=config.model.num_classes, anchors=config.model.anchors)
        model.classifier = ConvolutionalClassifier(num_classes=config.model.num_classes)
        model.to(self.device)

        return best_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)
        
       
        criterions['yolo_loss'] = YOLOLoss(num_classes=config.model.num_classes, anchors=config.model.anchors).to(self.device)
        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 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 train_step(self, batch, batch_idx):
        self.model.train()
        with autocast(device_type=self.device.type):
            images = batch['INPUT_IMAGE_KEY'].to(self.device)
            gt_boxes = batch['boxes'].to(self.device)  # Assuming 'boxes' exist
            gt_confidences = batch['INPUT_CONFIDENCE_KEY']  # Keep as list
            gt_classes = batch['TARGET_FISHING_KEY'].to(self.device)  # Tensor

            grid_size = 13  # Example grid size; adjust as needed

            # Encode targets using the box coder
            yolo_targets = self.box_coder.encode(
                centers=gt_boxes.cpu().numpy(),                # [N, 2] as NumPy array
                confidences=gt_confidences,                    # [N] as list of categorical labels
                lengths=batch['TARGET_LENGTHS_KEY'].cpu().numpy(),  # [N] as NumPy array
                is_vessel=batch['TARGET_VESSEL_KEY'].cpu().numpy(),   # [N] as NumPy array
                is_fishing=batch['TARGET_FISHING_KEY'].cpu().numpy()  # [N] as NumPy array
            )

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

            # Forward pass
            outputs = self.model(images)  # Expected output shape: [batch_size, 6, grid_size, grid_size]

            # Compute loss
            loss = self.criterion(outputs, {
                'heatmap': yolo_heatmap,
                'is_vessel': yolo_is_vessel,
                'is_fishing': yolo_is_fishing,
                'lengths': yolo_length,
                'offset': yolo_offset
            })  # Ensure `self.criterion` accepts a dictionary

        # 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 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)
        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(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
        )
        model = nas.search(config.get('nas', {}).get('num_iterations', 4), self.loaders['train'], self.loaders['valid'])

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

    def evaluate_architecture(self, model, train_loader, val_loader, device):
        """
        Compute the average validation loss over the validation DataLoader.
        """
        self.device = device
        model.to(self.device)
        model.eval()
        total_loss = 0.0
        total_samples = 0

        with torch.no_grad():
            for batch_idx, batch in enumerate(val_loader):
                # Debug: Print batch keys
                print(f"[DEBUG] Batch {batch_idx} keys: {batch.keys()}")

                # Check if 'INPUT_IMAGE_KEY' exists
                if 'INPUT_IMAGE_KEY' not in batch:
                    available_keys = list(batch.keys())
                    raise KeyError(f"'INPUT_IMAGE_KEY' key not found in batch {batch_idx}. Available keys: {available_keys}")

                images = batch['INPUT_IMAGE_KEY'].to(device)
                if 'boxes' in batch:
                    gt_boxes = batch['boxes'].to(device)  # Shape: [batch_size, max_num_objs, 4]
                else:
                    # Reconstruct boxes from centers and lengths
                    centers = batch.get('TARGET_CENTERS_KEY')  # Shape: [batch_size, max_num_objs, 2]
                    lengths = batch.get('TARGET_LENGTHS_KEY')  # Shape: [batch_size, max_num_objs]
                    if centers is None or lengths is None:
                        raise KeyError("Batch must contain 'boxes' or both 'TARGET_CENTERS_KEY' and 'TARGET_LENGTHS_KEY'")

                    # Assuming lengths correspond to width and height
                    widths = lengths  # Replace with actual representation if different
                    heights = lengths  # Replace as needed

                    # Convert centers to a PyTorch tensor if it's a list
                    if isinstance(centers, list):
                        centers = torch.tensor(centers, dtype=torch.float32)  # Choose appropriate dtype
                    # Before line 624 in evaluate_architecture
                    logger.debug(f"Type of centers: {type(centers)}")
                    logger.debug(f"centers: {centers}")
                    # Compute x_min, y_min, x_max, y_max
                    x_min = centers[..., 0] - (widths / 2)
                    y_min = centers[..., 1] - (heights / 2)
                    x_max = centers[..., 0] + (widths / 2)
                    y_max = centers[..., 1] + (heights / 2)

                    # Clamp coordinates to image boundaries (assuming image_size is available)
                    img_width, img_height = self.box_coder.image_size  # Define based on your dataset/config
                    x_min = x_min.clamp(min=0, max=img_width)
                    y_min = y_min.clamp(min=0, max=img_height)
                    x_max = x_max.clamp(min=0, max=img_width)
                    y_max = y_max.clamp(min=0, max=img_height)

                    # Stack to get boxes: [batch_size, max_num_objs, 4]
                    boxes = torch.stack([x_min, y_min, x_max, y_max], dim=3)

                    gt_boxes = boxes.to(device)

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

                centers = batch[TARGET_CENTERS_KEY].to(device)          # Now a tensor
                lengths = batch[TARGET_LENGTHS_KEY].to(device)          # Tensor
                is_vessel = batch[TARGET_VESSEL_KEY].to(device)        # Tensor
                is_fishing = batch[TARGET_FISHING_KEY].to(device)      # Tensor
                heatmap = batch[CENTERNET_TARGET_OBJECTNESS_MAP].to(device)  # Tensor
                vessel_map = batch[CENTERNET_TARGET_VESSEL_MAP].to(device)    # Tensor
                fishing_map = batch[CENTERNET_TARGET_FISHING_MAP].to(device)  # Tensor
                offset = batch[CENTERNET_TARGET_OFFSET].to(device)            # Tensor
                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'")

                gt_confidences = gt_confidences  # Already a list; no tensor conversion
                gt_classes = gt_classes.to(device)  # Tensor

                grid_size = 13  # Ensure consistency with training

                # Encode targets using the box coder
                yolo_targets = self.box_coder.encode(
                    centers=gt_boxes.cpu().numpy(),
                    confidences=gt_confidences,               # [N] as list of categorical labels
                    lengths=lengths.cpu().numpy(),           # [N] as NumPy array
                    is_vessel=is_vessel.cpu().numpy(),       # [N] as NumPy array
                    is_fishing=is_fishing.cpu().numpy()      # [N] 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)

                # Forward pass
                outputs = model(images)

                # Compute loss
                loss = self.criterion(outputs, {
                    'heatmap': yolo_heatmap,
                    'is_vessel': yolo_is_vessel,
                    'is_fishing': yolo_is_fishing,
                    'lengths': yolo_length,
                    'offset': yolo_offset
                })

                total_loss += loss.item() * images.size(0)
                total_samples += images.size(0)

        average_loss = total_loss / total_samples
        print(f"[DEBUG] Average Evaluation Loss: {average_loss}")

        return -average_loss  # Return negative loss if NAS maximizes the score
         
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()