that's why I wasn't converting confidence to jumpy as it's like this
detect_lat,detect_lon,vessel_length_m,source,detect_scene_row,detect_scene_column,is_vessel,is_fishing,distance_from_shore_km,scene_id,confidence,top,left,bottom,right,detect_id
5.66292355123629,4.84242888997755,,ais,16722,22703,,,9999.99,e42a50089e03990ft,LOW,,,,,e42a50089e03990ft_005.66292355123628965430_004.84242888997755027702
5.83078557395605,4.79439446548728,,ais,14867,22165,False,,9999.99,e42a50089e03990ft,HIGH,,,,,e42a50089e03990ft_005.83078557395605034941_004.79439446548728032127
5.65028885294943,5.07650030801318,,ais,16853,25297,False,,4.1204849040804055,e42a50089e03990ft,HIGH,,,,,e42a50089e03990ft_005.65028885294942995188_005.07650030801317964091
5.86549528127106,4.93833517222189,,ais,14478,23758,False,,4.34404231490167,e42a50089e03990ft,HIGH,,,,,e42a50089e03990ft_005.86549528127105990194_004.93833517222188955742
DEBUG:rasterio._base:Sharing flag: 0
Traceback (most recent call last):
File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel3.py", line 966, in
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: 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 "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel3.py", line 962, in main
pipeline.train()
File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel3.py", line 679, in train
train_loader, valid_loader, _, dataset_callbacks = self.build_datasets(config)
File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel3.py", line 311, in build_datasets
sample_batch = next(iter(train_loader))
File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 630, in next
data = self._next_data()
File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 1344, in _next_data
return self._process_data(data)
File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 1370, in _process_data
data.reraise()
File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/_utils.py", line 706, in reraise
raise exception
ValueError: Caught ValueError in DataLoader worker process 0.
Original Traceback (most recent call last):
File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/utils/data/_utils/worker.py", line 309, in _worker_loop
data = fetcher.fetch(index) # type: ignore[possibly-undefined]
File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 52, in fetch
data = [self.dataset[idx] for idx in possibly_batched_index]
File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 52, in
data = [self.dataset[idx] for idx in possibly_batched_index]
File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/dataset/random_crop_dataset.py", line 191, in getitem
return self.build_sample(
File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/multilabel_dataset.py", line 415, in build_sample
return self.build_centernet_sample(
File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/multilabel_dataset.py", line 222, in build_centernet_sample
confidences_np = np.array(confidences, dtype=np.float32)
ValueError: could not convert string to float: 'MEDIUM'

Thank you for providing the additional information and the error traceback. The issue you're encountering is due to the fact that the confidences field in your data contains string labels like 'LOW', 'MEDIUM', 'HIGH', and 'NONE', which cannot be directly converted to floats.
In the earlier advice, I suggested removing INPUT_CONFIDENCE_KEY from skip_keys and converting confidences to tensors, assuming they were numerical. However, since confidences are categorical string labels, we need to handle them differently.
Here's how we can fix the issue:
Keep INPUT_CONFIDENCE_KEY in skip_keys: Since confidences are categorical strings, we should not include them in the default collation process, which expects numerical tensors. By keeping INPUT_CONFIDENCE_KEY in skip_keys, we ensure that confidences are excluded from default_collate and remain as lists of strings. Revised centernet_collate function:

def centernet_collate(batch, channels_last=False):
    skip_keys = [
        TARGET_CENTERS_KEY,
        TARGET_FISHING_KEY,
        TARGET_VESSEL_KEY,
        TARGET_LABELS_KEY,
        TARGET_LENGTHS_KEY,
        INPUT_SCENE_ID_KEY,
        INPUT_SCENE_CROP_KEY,
        INPUT_CONFIDENCE_KEY,  # Keep this in skip_keys
        INPUT_LOCATION_ID_KEY,
        INPUT_IS_NEAR_SHORE_KEY,
    ]

    excluded_items = {k: [] for k in skip_keys}
    included_items = []

    for b in batch:
        included = {k: v for k, v in b.items() if k not in skip_keys}
        included_items.append(included)
        for k in skip_keys:
            if k in b:
                excluded_items[k].append(b[k])

    # Use default_collate for included items (numerical tensors)
    batch_dict = default_collate(included_items)

    # Process excluded items (lists of strings or other categorical data)
    for k, v in excluded_items.items():
        if v:
            batch_dict[k] = v  # Keep as a list of lists

    if channels_last and INPUT_IMAGE_KEY in batch_dict:
        batch_dict[INPUT_IMAGE_KEY] = batch_dict[INPUT_IMAGE_KEY].to(memory_format=torch.channels_last)

    return batch_dict
Avoid Converting confidences to np.float32 in build_centernet_sample: Since confidences are strings, we should not attempt to convert them directly to np.float32. Instead, we can keep them as lists of strings. When we need to use numerical values for confidences, we'll map them using self.confidence_mapping. Revised build_centernet_sample method:

def build_centernet_sample(
    self,
    image,
    centers,
    confidences,
    lengths,
    is_vessel,
    is_fishing,
    index,
    image_id,
    is_near_shore,
    location,
    folder: str,
    crop_coords: Tuple[Tuple[int, int], Tuple[int, int]],
) -> Dict[str, Any]:
    # Mask NaN values
    image[~np.isfinite(image)] = 0

    # Convert numerical fields to NumPy arrays with explicit dtypes
    centers_np = np.array(centers, dtype=np.float32)
    lengths_np = np.array(lengths, dtype=np.float32)
    is_vessel_np = np.array(is_vessel, dtype=np.float32)
    is_fishing_np = np.array(is_fishing, dtype=np.float32)
    is_near_shore_np = np.array(is_near_shore, dtype=np.float32)

    # No conversion for confidences here
    # confidences remain as a list of strings

    result = {
        INPUT_INDEX_KEY: index,
        INPUT_IMAGE_KEY: image_to_tensor(image),
        TARGET_CENTERS_KEY: centers_np,
        TARGET_VESSEL_KEY: is_vessel_np,
        TARGET_FISHING_KEY: is_fishing_np,
        TARGET_LENGTHS_KEY: lengths_np,
        INPUT_CONFIDENCE_KEY: confidences,  # Keep as list of strings
        INPUT_IS_NEAR_SHORE_KEY: is_near_shore_np,
        INPUT_LOCATION_ID_KEY: location,
        INPUT_FOLDER_KEY: folder,
        INPUT_SCENE_ID_KEY: image_id,
        TARGET_SAMPLE_WEIGHT: {"train": 0.25, "validation": 1.0, "valid": 1.0, "test": 0.1}[folder],
    }
    if crop_coords is not None:
        result[INPUT_SCENE_CROP_KEY] = crop_coords

    # Encode targets using the box coder
    box_coder = self.box_coder.create_box_coder_for_image_size(image.shape[:2])
    centernet_targets: MultilabelCircleNetEncodeResult = box_coder.encode(
        centers=centers_np,
        confidences=confidences,  # Pass as list of strings
        lengths=lengths_np,
        is_vessel=is_vessel_np,
        is_fishing=is_fishing_np,
    )

    # Convert encoded targets to PyTorch tensors
    result[CENTERNET_TARGET_OBJECTNESS_MAP] = torch.from_numpy(centernet_targets.heatmap).float()
    result[CENTERNET_TARGET_VESSEL_MAP] = torch.from_numpy(centernet_targets.is_vessel).float()
    result[CENTERNET_TARGET_FISHING_MAP] = torch.from_numpy(centernet_targets.is_fishing).float()
    result[CENTERNET_TARGET_SIZE] = torch.from_numpy(centernet_targets.lengths).float()
    result[CENTERNET_TARGET_OFFSET] = torch.from_numpy(centernet_targets.offset).float()
    result[CENTERNET_TARGET_YOLO] = torch.from_numpy(centernet_targets.yolo_targets).float()
    # Assuming centernet_targets has a yolo_targets attribute

    return result
Handle confidences in generate_yolo_targets Method: When generating YOLO targets, we'll map the string labels of confidences to numerical values using self.confidence_mapping. This way, we can work with numerical values during target generation without needing to convert the entire list at once. Revised generate_yolo_targets method:

def generate_yolo_targets(self, centers, confidences, is_vessel, is_fishing, lengths):
    """
    Generate YOLO-style targets with 6 channels: [x, y, w, h, confidence, class].
    """
    num_objects = len(centers)
    yolo_targets = np.zeros((self.grid_height, self.grid_width, 6), dtype=np.float32)

    for i in range(num_objects):
        center = centers[i]
        confidence_label = confidences[i]
        confidence = self.confidence_mapping[confidence_label]  # Map string label to numerical value
        vessel_type = self.vessel_mapping.get(is_vessel[i], self.ignore_value)
        fishing_status = self.fishing_mapping.get(is_fishing[i], self.ignore_value)
        length = lengths[i]

        # Compute grid indices
        grid_x = int(center[0] * self.grid_width)
        grid_y = int(center[1] * self.grid_height)

        # Ensure grid indices are within bounds
        grid_x = min(max(grid_x, 0), self.grid_width - 1)
        grid_y = min(max(grid_y, 0), self.grid_height - 1)

        # Assign the values
        yolo_targets[grid_y, grid_x, 0:4] = np.array([center[0], center[1], length, length])
        yolo_targets[grid_y, grid_x, 4] = confidence
        yolo_targets[grid_y, grid_x, 5] = vessel_type

    return yolo_targets
Ensure that self.confidence_mapping is a dictionary mapping string labels to numerical values, for example:

self.confidence_mapping = {
    'NONE': 0.0,
    'LOW': 0.33,
    'MEDIUM': 0.66,
    'HIGH': 1.0,
}
Ensure box_coder.encode Method Handles confidences Appropriately: In your MultilabelCircleNetCoder class, the encode method should be designed to work with confidences as lists of strings and perform the necessary mapping within the method. Sample implementation:

class MultilabelCircleNetCoder:
    # Existing methods...

    def encode(self, centers, confidences, lengths, is_vessel, is_fishing):
        # Map confidences from strings to numerical values
        confidences_numeric = np.array([self.confidence_mapping[c] for c in confidences], dtype=np.float32)

        # Proceed with encoding using numerical confidences
        # ...

        # Return an instance of MultilabelCircleNetEncodeResult
        return MultilabelCircleNetEncodeResult(
            heatmap=heatmap,
            is_vessel=is_vessel_map,
            is_fishing=is_fishing_map,
            lengths=lengths_map,
            offset=offset_map,
            yolo_targets=yolo_targets,  # Add yolo_targets if needed
        )
Update the MultilabelCircleNetTargetsMixin Initialization: Ensure that confidence_mapping is properly initialized in your mixin class. Updated __init__ method:

class MultilabelCircleNetTargetsMixin:
    def __init__(
        self,
        stride=32,
        confidence_mapping=None,
        vessel_mapping=None,
        fishing_mapping=None,
        ignore_value=-1,
    ):
        self.stride = stride
        self.confidence_mapping = confidence_mapping or {
            'NONE': 0.0,
            'LOW': 0.33,
            'MEDIUM': 0.66,
            'HIGH': 1.0,
        }
        self.vessel_mapping = vessel_mapping or {False: 0, True: 1}
        self.fishing_mapping = fishing_mapping or {False: 0, True: 1}
        self.ignore_value = ignore_value
Verify Data Consistency in build_centernet_sample: Since confidences is no longer being converted to a NumPy array, adjust the consistency check accordingly. Adjusted consistency check:

num_objects = len(confidences)
if not all([
    len(centers_np) == num_objects,
    len(lengths_np) == num_objects,
    len(is_vessel_np) == num_objects,
    len(is_fishing_np) == num_objects,
]):
    raise ValueError("Mismatch in number of objects across targets")
Remove Conversion of confidences to NumPy Array in build_centernet_sample: As we've adjusted the code to handle confidences as lists of strings, we can safely remove the line that was causing the error.

# Remove or comment out this line
# confidences_np = np.array(confidences, dtype=np.float32)  # This line causes the ValueError
By following these adjustments, your code should now correctly handle the confidences field containing string labels, mapping them to numerical values only when needed. This approach maintains the integrity of your data and ensures that the confidences are appropriately processed during target generation.
Key Points to Remember:
Categorical Data Handling: For categorical variables represented as strings, avoid converting them directly to numerical types without mapping. Keep them as strings unless numerical values are required.
Mapping at Point of Use: Map categorical strings to numerical values only when needed, such as during target generation or computation.
Collation of Categorical Data: Exclude categorical data from default collation by including their keys in skip_keys. Process them separately to maintain their structure (e.g., lists of strings).
Data Consistency: Always check for consistency in the number of objects across all target arrays or lists to prevent misalignment issues.
Testing Your Changes:
Unit Testing: Test your build_centernet_sample and generate_yolo_targets methods independently with sample data to ensure they handle confidences correctly.
Small Batch Testing: Run a small batch through your data loader to verify that the collate function and data processing steps work without errors.
Debugging: If errors persist, add debug statements to print the types and values of variables at critical points in your code to trace the source of issues.
Example Usage:

# Example data
centers = [[0.5, 0.5], [0.25, 0.75]]
confidences = ['HIGH', 'LOW']
is_vessel = [True, False]
is_fishing = [False, True]
lengths = [10.0, 15.0]

# Instantiate the coder
box_coder = MultilabelCircleNetCoder()
targets_mixin = MultilabelCircleNetTargetsMixin()

# Build sample
sample = targets_mixin.build_centernet_sample(
    image=np.random.rand(512, 512, 3),
    centers=centers,
    confidences=confidences,
    lengths=lengths,
    is_vessel=is_vessel,
    is_fishing=is_fishing,
    index=0,
    image_id='image_001',
    is_near_shore=[True, False],
    location='location_1',
    folder='train',
    crop_coords=None,
)

# Check the sample
print(sample[INPUT_CONFIDENCE_KEY])  # Should print ['HIGH', 'LOW']
print(sample[CENTERNET_TARGET_OBJECTNESS_MAP].shape)
By carefully handling the confidences field and ensuring that mappings are applied at the correct stages, you should be able to integrate your YOLO-style target generation without encountering the ValueError related to string-to-float conversion.
If you have any further questions or need assistance with specific parts of your code, feel free to ask!

I think there's critical issues and many mistakes I have made when I was trying to integrate yolo in the multilabel_circle_coder.py and also multilabel_trimmed.loss.py files!! and consequently the multilabel_dataset.py so to fix let's see
the original
###multilabel_trimmed.loss.py
from functools import partial
from typing import Dict, Tuple
import torch
from pytorch_toolbelt.losses import SoftBCEWithLogitsLoss, BinaryBiTemperedLogisticLoss, BinaryFocalLoss, BalancedBCEWithLogitsLoss
from torch import nn, Tensor
from torch.nn import functional as F
from .functional import binary_shannon_entropy_loss
from ..constants import (
CENTERNET_OUTPUT_SIZE,
CENTERNET_OUTPUT_OFFSET,
CENTERNET_OUTPUT_OBJECTNESS_MAP,
CENTERNET_TARGET_SIZE,
CENTERNET_TARGET_OFFSET,
CENTERNET_TARGET_FISHING_MAP,
CENTERNET_OUTPUT_FISHING_MAP,
CENTERNET_OUTPUT_VESSEL_MAP,
CENTERNET_TARGET_VESSEL_MAP,
CENTERNET_TARGET_OBJECTNESS_MAP,
)
from ...constants import IGNORE_LABEL
all = ["MultilabelCenterTrimmedLoss"]
class MultilabelCenterTrimmedLoss(nn.Module):
fraction: int
samplewise: bool
regularize_ignored: bool

def __init__(
    self,
    fraction: int,
    samplewise: bool = True,
    objectness_loss: str = "rfl",
    classifier_loss="bce",
    size_loss: str = "mse",
    offset_loss: str = "mse",
    regularize_ignored: bool = False,
    regularization_weight: float = 1.0,
    ignore_index: int = IGNORE_LABEL,
):
    super().__init__()
    self.fraction = fraction
    self.objectness_loss = objectness_loss
    self.samplewise = samplewise
    self.regularize_ignored = regularize_ignored
    self.regularization_weight = regularization_weight
    self.ignore_index = ignore_index

    if objectness_loss == "rfl":
        self.objectness_loss = partial(self.reduced_focal_loss, alpha=2, beta=4, eps=0)
    elif objectness_loss == "l1":
        self.objectness_loss = self.objectness_l1_loss
    elif objectness_loss == "mse":
        self.objectness_loss = self.objectness_mse_loss
    elif objectness_loss == "focal_mse":
        self.objectness_loss = self.objectness_focal_mse_loss
    else:
        raise KeyError(objectness_loss)

    if classifier_loss == "bce":
        self.classifier_loss = SoftBCEWithLogitsLoss(ignore_index=ignore_index, reduction="none")
    elif classifier_loss == "bbce":
        self.classifier_loss = BalancedBCEWithLogitsLoss(ignore_index=ignore_index, reduction="none")
    elif classifier_loss == "soft_bce":
        self.classifier_loss = SoftBCEWithLogitsLoss(ignore_index=ignore_index, reduction="none", smooth_factor=0.05)
    elif classifier_loss == "bitempered":
        self.classifier_loss = BinaryBiTemperedLogisticLoss(t1=0.5, t2=4.0, ignore_index=ignore_index, reduction="none")
    elif classifier_loss == "focal":
        self.classifier_loss = BinaryFocalLoss(ignore_index=ignore_index, alpha=None, reduction="none")
    else:
        raise KeyError(classifier_loss)

    if size_loss == "huber":
        self.size_loss = self.huber_loss
    elif size_loss == "mse":
        self.size_loss = self.mse_loss
    elif size_loss == "smooth_l1":
        self.size_loss = self.smooth_l1_loss
    elif size_loss == "rel_error":
        self.size_loss = self.relative_length_error
    else:
        raise KeyError(size_loss)

    if offset_loss == "mse":
        self.offset_loss = self.mse_loss
    elif offset_loss == "l1":
        self.offset_loss = self.l1_loss
    elif offset_loss == "smooth_l1":
        self.offset_loss = self.smooth_l1_loss
    elif offset_loss == "huber":
        self.offset_loss = self.huber_loss
    elif offset_loss is None:
        self.offset_loss = None
    else:
        raise KeyError(offset_loss)

def objectness_mse_loss(self, pred: Tensor, gt: Tensor, pos_mask) -> Tuple[Tensor, Tensor]:
    ign_mask = gt.eq(self.ignore_index)
    pt = pred.sigmoid()
    neg_mask = ~pos_mask
    loss = F.mse_loss(pt, gt, reduction="none")
    pos_loss: Tensor = loss * pos_mask
    neg_loss: Tensor = loss * neg_mask
    return torch.masked_fill(pos_loss, ign_mask, 0), torch.masked_fill(neg_loss, ign_mask, 0)

def objectness_focal_mse_loss(self, pred: Tensor, gt: Tensor, pos_mask: Tensor, alpha=2.0, beta=2.0) -> Tuple[Tensor, Tensor]:
    pt = pred.sigmoid()
    neg_mask = ~pos_mask
    loss = F.mse_loss(pt, gt, reduction="none")
    pos_loss: Tensor = loss * pos_mask
    neg_loss: Tensor = loss * neg_mask * torch.pow(pt, alpha) * torch.pow(1.0 - gt, beta)
    return pos_loss, neg_loss

def objectness_l1_loss(self, pred: Tensor, gt: Tensor, pos_mask):
    pt = pred.sigmoid()
    ign_mask = gt.eq(self.ignore_index)
    neg_mask = ~pos_mask
    loss = F.l1_loss(pt, gt, reduction="none")
    pos_loss: Tensor = loss * pos_mask
    neg_loss: Tensor = (loss * neg_mask) ** 4
    return torch.masked_fill(pos_loss, ign_mask, 0), torch.masked_fill(neg_loss, ign_mask, 0)

def reduced_focal_loss(self, pred: Tensor, gt: Tensor, pos_mask, alpha: float, beta: float, eps: float = 1e-4):
    neg_mask = ~pos_mask
    ign_mask = gt.eq(self.ignore_index)

    pt = pred.sigmoid().clamp(eps, 1 - eps)

    pos_loss: Tensor = -torch.pow(1 - pt, alpha) * F.logsigmoid(pred) * pos_mask
    neg_loss: Tensor = -torch.pow(1.0 - gt, beta) * torch.pow(pt, alpha) * F.logsigmoid(-pred) * neg_mask

    return torch.masked_fill(pos_loss, ign_mask, 0), torch.masked_fill(neg_loss, ign_mask, 0)

def mse_loss(self, prediction, target, pos_mask):
    loss = F.mse_loss(prediction, target, reduction="none") * pos_mask
    return loss

def l1_loss(self, prediction, target, pos_mask):
    loss = F.l1_loss(prediction, target, reduction="none") * pos_mask
    return loss

def smooth_l1_loss(self, prediction, target, pos_mask):
    loss = F.smooth_l1_loss(prediction, target, reduction="none") * pos_mask
    return loss

def huber_loss(self, prediction, target, pos_mask):
    loss = F.huber_loss(prediction, target, reduction="none", delta=1) * pos_mask
    return loss

def relative_length_error(self, prediction, target, pos_mask, eps=1e-4):
    loss = pos_mask * F.smooth_l1_loss(prediction, target, reduction="none") / (target + eps)
    return loss

# @torch.cuda.amp.autocast(False)
def forward(self, **kwargs) -> Dict[str, Tensor]:
    gt_objectness: Tensor = kwargs[CENTERNET_TARGET_OBJECTNESS_MAP]
    gt_is_vessel: Tensor = kwargs[CENTERNET_TARGET_VESSEL_MAP]
    gt_is_fishing: Tensor = kwargs[CENTERNET_TARGET_FISHING_MAP]
    gt_size: Tensor = kwargs[CENTERNET_TARGET_SIZE]

    pred_objectness: Tensor = kwargs[CENTERNET_OUTPUT_OBJECTNESS_MAP]
    pred_is_vessel: Tensor = kwargs[CENTERNET_OUTPUT_VESSEL_MAP]
    pred_is_fishing: Tensor = kwargs[CENTERNET_OUTPUT_FISHING_MAP]
    pred_size: Tensor = kwargs[CENTERNET_OUTPUT_SIZE]

    peaks_mask = gt_objectness.eq(1)
    true_vessels_mask = gt_is_vessel.eq(1)

    num_objects = peaks_mask.sum(dtype=torch.float32).clamp_min(1.0)
    num_vessels = (true_vessels_mask & peaks_mask).sum(dtype=torch.float32).clamp_min(1.0)

    uncertain_objectness_mask = gt_objectness.eq(self.ignore_index)
    certain_objectness_mask = ~uncertain_objectness_mask

    # Returns unreduced losses
    pos_objectness_loss, neg_objectness_loss = self.objectness_loss(pred_objectness, gt_objectness, pos_mask=peaks_mask)

    # Weight classification loss by objectness (This make sense for target encoding=center).
    # This way pixels closer to the center of the keypoint will have bigger influence than ones on the periphery
    vessel_loss = self.classifier_loss(pred_is_vessel, gt_is_vessel) * gt_objectness * certain_objectness_mask
    fishing_loss = self.classifier_loss(pred_is_fishing, gt_is_fishing) * gt_objectness * certain_objectness_mask

    if self.regularize_ignored:
        uncertain_vessel_mask = gt_is_vessel.eq(self.ignore_index)
        uncertain_fishing_mask = gt_is_fishing.eq(self.ignore_index)
        # Plus regularization loss for ignored values in the objectness heatmap
        objectness_reg_loss = binary_shannon_entropy_loss(pred_objectness) * uncertain_objectness_mask

        vsl_reg_loss = binary_shannon_entropy_loss(pred_is_vessel) * certain_objectness_mask * uncertain_vessel_mask * gt_objectness
        fsh_reg_loss = (
            binary_shannon_entropy_loss(pred_is_fishing)
            * certain_objectness_mask
            * uncertain_fishing_mask
            * true_vessels_mask
            * gt_objectness
        )

        pos_objectness_loss += objectness_reg_loss * self.regularization_weight
        vessel_loss += vsl_reg_loss * self.regularization_weight
        fishing_loss += fsh_reg_loss * self.regularization_weight

    valid_size_mask = peaks_mask * (gt_size > 0)  # Size loss only for positive lengths
    num_objects_with_length = valid_size_mask.sum(dtype=torch.float32).clamp_min(1.0)
    size_loss = self.size_loss(pred_size, gt_size, valid_size_mask)

    if self.offset_loss is not None:
        offset_loss = self.offset_loss(kwargs[CENTERNET_OUTPUT_OFFSET], kwargs[CENTERNET_TARGET_OFFSET], peaks_mask).sum(
            dim=1, keepdim=True
        )
    else:
        offset_loss = torch.zeros_like(size_loss)

    if self.fraction > 0:
        if self.samplewise:
            (pos_objectness_loss, neg_objectness_loss, vessel_loss, fishing_loss, offset_loss, size_loss,) = samplewise_trimming(
                pred_objectness,
                pred_is_vessel,
                pred_is_fishing,
                pos_objectness_loss,
                neg_objectness_loss,
                vessel_loss,
                fishing_loss,
                offset_loss,
                size_loss,
            )
        else:
            (pos_objectness_loss, neg_objectness_loss, vessel_loss, fishing_loss, offset_loss, size_loss,) = batchwise_trimming(
                pred_objectness,
                pred_is_vessel,
                pred_is_fishing,
                pos_objectness_loss,
                neg_objectness_loss,
                vessel_loss,
                fishing_loss,
                offset_loss,
                size_loss,
            )

    mean_obj_loss = (neg_objectness_loss.sum(dtype=torch.float32) + pos_objectness_loss.sum(dtype=torch.float32)) / num_objects
    mean_vessel_loss = (vessel_loss.sum(dtype=torch.float32)) / num_objects
    mean_fishing_loss = (fishing_loss.sum(dtype=torch.float32)) / num_vessels
    mean_siz_loss = size_loss.sum(dtype=torch.float32) / num_objects_with_length

    losses = {
        "objectness": mean_obj_loss.type_as(pred_objectness),
        "vessel": mean_vessel_loss.type_as(pred_objectness),
        "fishing": mean_fishing_loss.type_as(pred_objectness),
        "size": mean_siz_loss.type_as(pred_objectness),
    }
    if self.offset_loss is not None:
        mean_off_loss = offset_loss.sum(dtype=torch.float32) / num_objects
        losses["offset"] = mean_off_loss.type_as(pred_objectness)
    return losses
def samplewise_trimming(
pred_objectness,
pred_is_vessel,
pred_is_fishing,
pos_objectness_loss,
neg_objectness_loss,
vessel_loss,
fishing_loss,
offset_loss,
size_loss,
num_elements_to_ignore,
):
pred_objectness = pred_objectness.flatten(2) # [B, 1, N]
pred_is_vessel = pred_is_vessel.flatten(2) # [B, 1, N]
pred_is_fishing = pred_is_fishing.flatten(2) # [B, 1, N]
pos_objectness_loss = pos_objectness_loss.flatten(2) # [B, 1, N]
neg_objectness_loss = neg_objectness_loss.flatten(2) # [B, 1, N]
vessel_loss = vessel_loss.flatten(2) # [B, 1, N]
fishing_loss = fishing_loss.flatten(2) # [B, 1, N]
offset_loss = offset_loss.flatten(2) # [B, 1, N]
size_loss = size_loss.flatten(2) # [B, 1, N]

# Here we ignore top-N negative losses for negative objectness loss (E.g target is predicted, but groundtruth is negative).
# Since some objects may be missing in gt, and we don't want to include those wrong false-positives in the training signal
# And use keep indexes to propagate ignored pixels in other losses
_, indices = torch.topk(neg_objectness_loss, k=num_elements_to_ignore, dim=2, largest=True, sorted=False)

non_ignored_mask = torch.ones_like(neg_objectness_loss).scatter_(2, indices, 0)

# Instead of just setting loss to zeros in those hard example predictions, we minimize entropy loss
# for objectness & classification to push the predictions to either 0 or 1
neg_objectness_loss = neg_objectness_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_objectness) * (1 - non_ignored_mask)
vessel_loss = vessel_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_is_vessel) * (1 - non_ignored_mask)
fishing_loss = fishing_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_is_fishing) * (1 - non_ignored_mask)
# For offset & size there is no regularization, so we exclude corresponding pixels from the grad. contribution
offset_loss = offset_loss * non_ignored_mask
size_loss = size_loss * non_ignored_mask

return (
    pos_objectness_loss,
    neg_objectness_loss,
    vessel_loss,
    fishing_loss,
    offset_loss,
    size_loss,
)
def batchwise_trimming(
pred_objectness,
pred_is_vessel,
pred_is_fishing,
pos_objectness_loss,
neg_objectness_loss,
vessel_loss,
fishing_loss,
offset_loss,
size_loss,
num_elements_to_ignore,
):
pred_objectness = pred_objectness.permute(1, 0, 2, 3).flatten(1) # [1, BxN]
pred_is_vessel = pred_is_vessel.permute(1, 0, 2, 3).flatten(1) # [C, BxN]
pred_is_fishing = pred_is_fishing.permute(1, 0, 2, 3).flatten(1) # [C, BxN]
pos_objectness_loss = pos_objectness_loss.permute(1, 0, 2, 3).flatten(1) # [1, BxN]
neg_objectness_loss = neg_objectness_loss.permute(1, 0, 2, 3).flatten(1) # [1, BxN]
vessel_loss = vessel_loss.permute(1, 0, 2, 3).flatten(1) # [1, BxN]
fishing_loss = fishing_loss.permute(1, 0, 2, 3).flatten(1) # [1, BxN]
offset_loss = offset_loss.permute(1, 0, 2, 3).flatten(1) # [1, BxN]
size_loss = size_loss.permute(1, 0, 2, 3).flatten(1) # [1, BxN]

# Here we ignore top-N negative losses for negative objectness loss (E.g target is predicted, but groundtruth is negative).
# Since some objects may be missing in gt, and we don't want to include those wrong false-positives in the training signal
# And use keep indexes to propagate ignored pixels in other losses

_, indices = torch.topk(neg_objectness_loss, k=num_elements_to_ignore, dim=1, largest=True, sorted=False)

non_ignored_mask = torch.ones_like(neg_objectness_loss).scatter_(1, indices, 0)

# Instead of just setting loss to zeros in those hard example predictions, we minimize entropy loss
# for objectness & classification to push the predictions to either 0 or 1
neg_objectness_loss = neg_objectness_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_objectness) * (1 - non_ignored_mask)
vessel_loss = vessel_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_is_vessel) * (1 - non_ignored_mask)
fishing_loss = fishing_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_is_fishing) * (1 - non_ignored_mask)
# For offset & size there is no regularization, so we exclude corresponding pixels from the grad. contribution
offset_loss = offset_loss * non_ignored_mask
size_loss = size_loss * non_ignored_mask

return (
    pos_objectness_loss,
    neg_objectness_loss,
    vessel_loss,
    fishing_loss,
    offset_loss,
    size_loss,
)
###my upgraded multilabel_trimmed.loss.py
from functools import partial
from typing import Dict, Tuple
import torch
from pytorch_toolbelt.losses import SoftBCEWithLogitsLoss, BinaryBiTemperedLogisticLoss, BinaryFocalLoss, BalancedBCEWithLogitsLoss
from torch import nn, Tensor
from torch.nn import functional as F
from .functional import binary_shannon_entropy_loss
from ..constants import (
CENTERNET_OUTPUT_SIZE,
CENTERNET_OUTPUT_OFFSET,
CENTERNET_OUTPUT_OBJECTNESS_MAP,
CENTERNET_TARGET_SIZE,
CENTERNET_TARGET_OFFSET,
CENTERNET_TARGET_FISHING_MAP,
CENTERNET_OUTPUT_FISHING_MAP,
CENTERNET_OUTPUT_VESSEL_MAP,
CENTERNET_TARGET_VESSEL_MAP,
CENTERNET_TARGET_OBJECTNESS_MAP,
CENTERNET_OUTPUT_YOLO,
CENTERNET_TARGET_YOLO
)
from ...constants import IGNORE_LABEL
all = ["MultilabelCenterTrimmedLoss"]
class MultilabelCenterTrimmedLoss(nn.Module):
fraction: int
samplewise: bool
regularize_ignored: bool

def __init__(
    self,
    fraction: int,
    samplewise: bool = True,
    objectness_loss: str = "rfl",
    classifier_loss="bce",
    size_loss: str = "mse",
    offset_loss: str = "mse",
    regularize_ignored: bool = False,
    regularization_weight: float = 1.0,
    ignore_index: int = IGNORE_LABEL,
    num_classes: int = 20,
    anchors: list = [[10, 13], [16, 30], [33, 23]],
):
    super().__init__()
    self.fraction = fraction
    self.objectness_loss = objectness_loss
    self.samplewise = samplewise
    self.regularize_ignored = regularize_ignored
    self.regularization_weight = regularization_weight
    self.ignore_index = ignore_index

    if objectness_loss == "rfl":
        self.objectness_loss = partial(self.reduced_focal_loss, alpha=2, beta=4, eps=0)
    elif objectness_loss == "l1":
        self.objectness_loss = self.objectness_l1_loss
    elif objectness_loss == "mse":
        self.objectness_loss = self.objectness_mse_loss
    elif objectness_loss == "focal_mse":
        self.objectness_loss = self.objectness_focal_mse_loss
    else:
        raise KeyError(objectness_loss)

    if classifier_loss == "bce":
        self.classifier_loss = SoftBCEWithLogitsLoss(ignore_index=ignore_index, reduction="none")
    elif classifier_loss == "bbce":
        self.classifier_loss = BalancedBCEWithLogitsLoss(ignore_index=ignore_index, reduction="none")
    elif classifier_loss == "soft_bce":
        self.classifier_loss = SoftBCEWithLogitsLoss(ignore_index=ignore_index, reduction="none", smooth_factor=0.05)
    elif classifier_loss == "bitempered":
        self.classifier_loss = BinaryBiTemperedLogisticLoss(t1=0.5, t2=4.0, ignore_index=ignore_index, reduction="none")
    elif classifier_loss == "focal":
        self.classifier_loss = BinaryFocalLoss(ignore_index=ignore_index, alpha=None, reduction="none")
    else:
        raise KeyError(classifier_loss)

    if size_loss == "huber":
        self.size_loss = self.huber_loss
    elif size_loss == "mse":
        self.size_loss = self.mse_loss
    elif size_loss == "smooth_l1":
        self.size_loss = self.smooth_l1_loss
    elif size_loss == "rel_error":
        self.size_loss = self.relative_length_error
    else:
        raise KeyError(size_loss)

    if offset_loss == "mse":
        self.offset_loss = self.mse_loss
    elif offset_loss == "l1":
        self.offset_loss = self.l1_loss
    elif offset_loss == "smooth_l1":
        self.offset_loss = self.smooth_l1_loss
    elif offset_loss == "huber":
        self.offset_loss = self.huber_loss
    elif offset_loss is None:
        self.offset_loss = None
    else:
        raise KeyError(offset_loss)

    # Initialize YOLO loss module
    self.yolo_loss_module = YOLOLoss(num_classes=num_classes, anchors=anchors)

def objectness_mse_loss(self, pred: Tensor, gt: Tensor, pos_mask) -> Tuple[Tensor, Tensor]:
    ign_mask = gt.eq(self.ignore_index)
    pt = pred.sigmoid()
    neg_mask = ~pos_mask
    loss = F.mse_loss(pt, gt, reduction="none")
    pos_loss: Tensor = loss * pos_mask
    neg_loss: Tensor = loss * neg_mask
    return torch.masked_fill(pos_loss, ign_mask, 0), torch.masked_fill(neg_loss, ign_mask, 0)

def objectness_focal_mse_loss(self, pred: Tensor, gt: Tensor, pos_mask: Tensor, alpha=2.0, beta=2.0) -> Tuple[Tensor, Tensor]:
    pt = pred.sigmoid()
    neg_mask = ~pos_mask
    loss = F.mse_loss(pt, gt, reduction="none")
    pos_loss: Tensor = loss * pos_mask
    neg_loss: Tensor = loss * neg_mask * torch.pow(pt, alpha) * torch.pow(1.0 - gt, beta)
    return pos_loss, neg_loss

def objectness_l1_loss(self, pred: Tensor, gt: Tensor, pos_mask):
    pt = pred.sigmoid()
    ign_mask = gt.eq(self.ignore_index)
    neg_mask = ~pos_mask
    loss = F.l1_loss(pt, gt, reduction="none")
    pos_loss: Tensor = loss * pos_mask
    neg_loss: Tensor = (loss * neg_mask) ** 4
    return torch.masked_fill(pos_loss, ign_mask, 0), torch.masked_fill(neg_loss, ign_mask, 0)

def reduced_focal_loss(self, pred: Tensor, gt: Tensor, pos_mask, alpha: float, beta: float, eps: float = 1e-4):
    neg_mask = ~pos_mask
    ign_mask = gt.eq(self.ignore_index)

    pt = pred.sigmoid().clamp(eps, 1 - eps)

    pos_loss: Tensor = -torch.pow(1 - pt, alpha) * F.logsigmoid(pred) * pos_mask
    neg_loss: Tensor = -torch.pow(1.0 - gt, beta) * torch.pow(pt, alpha) * F.logsigmoid(-pred) * neg_mask

    return torch.masked_fill(pos_loss, ign_mask, 0), torch.masked_fill(neg_loss, ign_mask, 0)

def mse_loss(self, prediction, target, pos_mask):
    loss = F.mse_loss(prediction, target, reduction="none") * pos_mask
    return loss

def l1_loss(self, prediction, target, pos_mask):
    loss = F.l1_loss(prediction, target, reduction="none") * pos_mask
    return loss

def smooth_l1_loss(self, prediction, target, pos_mask):
    loss = F.smooth_l1_loss(prediction, target, reduction="none") * pos_mask
    return loss

def huber_loss(self, prediction, target, pos_mask):
    loss = F.huber_loss(prediction, target, reduction="none", delta=1) * pos_mask
    return loss

def relative_length_error(self, prediction, target, pos_mask, eps=1e-4):
    loss = pos_mask * F.smooth_l1_loss(prediction, target, reduction="none") / (target + eps)
    return loss

# @torch.cuda.amp.autocast(False)
def forward(self, **kwargs) -> Dict[str, Tensor]:
    gt_objectness: Tensor = kwargs[CENTERNET_TARGET_OBJECTNESS_MAP]
    gt_is_vessel: Tensor = kwargs[CENTERNET_TARGET_VESSEL_MAP]
    gt_is_fishing: Tensor = kwargs[CENTERNET_TARGET_FISHING_MAP]
    gt_size: Tensor = kwargs[CENTERNET_TARGET_SIZE]
    gt_yolo: Tensor = kwargs(CENTERNET_TARGET_YOLO)
    if gt_yolo is None:
        raise ValueError(f"Missing key '{CENTERNET_TARGET_YOLO}' in input dictionary")

    pred_objectness: Tensor = kwargs[CENTERNET_OUTPUT_OBJECTNESS_MAP]
    pred_is_vessel: Tensor = kwargs[CENTERNET_OUTPUT_VESSEL_MAP]
    pred_is_fishing: Tensor = kwargs[CENTERNET_OUTPUT_FISHING_MAP]
    pred_size: Tensor = kwargs[CENTERNET_OUTPUT_SIZE]
    pred_yolo: Tensor = kwargs(CENTERNET_OUTPUT_YOLO, torch.zeros_like(pred_objectness))  # Default to zeros if not present

    peaks_mask = gt_objectness.eq(1)
    true_vessels_mask = gt_is_vessel.eq(1)

    num_objects = peaks_mask.sum(dtype=torch.float32).clamp_min(1.0)
    num_vessels = (true_vessels_mask & peaks_mask).sum(dtype=torch.float32).clamp_min(1.0)

    uncertain_objectness_mask = gt_objectness.eq(self.ignore_index)
    certain_objectness_mask = ~uncertain_objectness_mask

    # Returns unreduced losses
    pos_objectness_loss, neg_objectness_loss = self.objectness_loss(pred_objectness, gt_objectness, pos_mask=peaks_mask)

    vessel_loss = self.classifier_loss(pred_is_vessel, gt_is_vessel) * gt_objectness * certain_objectness_mask
    fishing_loss = self.classifier_loss(pred_is_fishing, gt_is_fishing) * gt_objectness * certain_objectness_mask
    yolo_loss = self.yolo_loss_module(pred_yolo, gt_yolo)

    if self.regularize_ignored:
        uncertain_vessel_mask = gt_is_vessel.eq(self.ignore_index)
        uncertain_fishing_mask = gt_is_fishing.eq(self.ignore_index)
        objectness_reg_loss = binary_shannon_entropy_loss(pred_objectness) * uncertain_objectness_mask

        vsl_reg_loss = binary_shannon_entropy_loss(pred_is_vessel) * certain_objectness_mask * uncertain_vessel_mask * gt_objectness
        fsh_reg_loss = (
            binary_shannon_entropy_loss(pred_is_fishing)
            * certain_objectness_mask
            * uncertain_fishing_mask
            * true_vessels_mask
            * gt_objectness
        )

        pos_objectness_loss += objectness_reg_loss * self.regularization_weight
        vessel_loss += vsl_reg_loss * self.regularization_weight
        fishing_loss += fsh_reg_loss * self.regularization_weight

    valid_size_mask = peaks_mask * (gt_size > 0)  # Size loss only for positive lengths
    num_objects_with_length = valid_size_mask.sum(dtype=torch.float32).clamp_min(1.0)
    size_loss = self.size_loss(pred_size, gt_size, valid_size_mask)

    if self.offset_loss is not None:
        offset_loss = self.offset_loss(kwargs[CENTERNET_OUTPUT_OFFSET], kwargs[CENTERNET_TARGET_OFFSET], peaks_mask).sum(
            dim=1, keepdim=True
        )
    else:
        offset_loss = torch.zeros_like(size_loss)

    if self.fraction > 0:
        if self.samplewise:
            (pos_objectness_loss, neg_objectness_loss, vessel_loss, fishing_loss, offset_loss, size_loss,) = samplewise_trimming(
                pred_objectness,
                pred_is_vessel,
                pred_is_fishing,
                pos_objectness_loss,
                neg_objectness_loss,
                vessel_loss,
                fishing_loss,
                offset_loss,
                size_loss,
            )
        else:
            (pos_objectness_loss, neg_objectness_loss, vessel_loss, fishing_loss, offset_loss, size_loss,) = batchwise_trimming(
                pred_objectness,
                pred_is_vessel,
                pred_is_fishing,
                pos_objectness_loss,
                neg_objectness_loss,
                vessel_loss,
                fishing_loss,
                offset_loss,
                size_loss,
            )

    mean_obj_loss = (neg_objectness_loss.sum(dtype=torch.float32) + pos_objectness_loss.sum(dtype=torch.float32)) / num_objects
    mean_vessel_loss = (vessel_loss.sum(dtype=torch.float32)) / num_objects
    mean_fishing_loss = (fishing_loss.sum(dtype=torch.float32)) / num_vessels
    mean_siz_loss = size_loss.sum(dtype=torch.float32) / num_objects_with_length

    losses = {
        "objectness": mean_obj_loss.type_as(pred_objectness),
        "vessel": mean_vessel_loss.type_as(pred_objectness),
        "fishing": mean_fishing_loss.type_as(pred_objectness),
        "size": mean_siz_loss.type_as(pred_objectness),
        "yolo_loss": yolo_loss.type_as(pred_objectness),  # Add YOLO loss to the losses dictionary
    }
    if self.offset_loss is not None:
        mean_off_loss = offset_loss.sum(dtype=torch.float32) / num_objects
        losses["offset"] = mean_off_loss.type_as(pred_objectness)
    return losses
def samplewise_trimming(
pred_objectness,
pred_is_vessel,
pred_is_fishing,
pos_objectness_loss,
neg_objectness_loss,
vessel_loss,
fishing_loss,
offset_loss,
size_loss,
num_elements_to_ignore,
):
pred_objectness = pred_objectness.flatten(2) # [B, 1, N]
pred_is_vessel = pred_is_vessel.flatten(2) # [B, 1, N]
pred_is_fishing = pred_is_fishing.flatten(2) # [B, 1, N]
pos_objectness_loss = pos_objectness_loss.flatten(2) # [B, 1, N]
neg_objectness_loss = neg_objectness_loss.flatten(2) # [B, 1, N]
vessel_loss = vessel_loss.flatten(2) # [B, 1, N]
fishing_loss = fishing_loss.flatten(2) # [B, 1, N]
offset_loss = offset_loss.flatten(2) # [B, 1, N]
size_loss = size_loss.flatten(2) # [B, 1, N]
_, indices = torch.topk(neg_objectness_loss, k=num_elements_to_ignore, dim=2, largest=True, sorted=False)

non_ignored_mask = torch.ones_like(neg_objectness_loss).scatter_(2, indices, 0)

neg_objectness_loss = neg_objectness_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_objectness) * (1 - non_ignored_mask)
vessel_loss = vessel_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_is_vessel) * (1 - non_ignored_mask)
fishing_loss = fishing_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_is_fishing) * (1 - non_ignored_mask)
offset_loss = offset_loss * non_ignored_mask
size_loss = size_loss * non_ignored_mask

return (
    pos_objectness_loss,
    neg_objectness_loss,
    vessel_loss,
    fishing_loss,
    offset_loss,
    size_loss,
)
def batchwise_trimming(
pred_objectness,
pred_is_vessel,
pred_is_fishing,
pos_objectness_loss,
neg_objectness_loss,
vessel_loss,
fishing_loss,
offset_loss,
size_loss,
num_elements_to_ignore,
):
pred_objectness = pred_objectness.permute(1, 0, 2, 3).flatten(1) # [1, BxN]
pred_is_vessel = pred_is_vessel.permute(1, 0, 2, 3).flatten(1) # [C, BxN]
pred_is_fishing = pred_is_fishing.permute(1, 0, 2, 3).flatten(1) # [C, BxN]
pos_objectness_loss = pos_objectness_loss.permute(1, 0, 2, 3).flatten(1) # [1, BxN]
neg_objectness_loss = neg_objectness_loss.permute(1, 0, 2, 3).flatten(1) # [1, BxN]
vessel_loss = vessel_loss.permute(1, 0, 2, 3).flatten(1) # [1, BxN]
fishing_loss = fishing_loss.permute(1, 0, 2, 3).flatten(1) # [1, BxN]
offset_loss = offset_loss.permute(1, 0, 2, 3).flatten(1) # [1, BxN]
size_loss = size_loss.permute(1, 0, 2, 3).flatten(1) # [1, BxN]

_, indices = torch.topk(neg_objectness_loss, k=num_elements_to_ignore, dim=1, largest=True, sorted=False)

non_ignored_mask = torch.ones_like(neg_objectness_loss).scatter_(1, indices, 0)

neg_objectness_loss = neg_objectness_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_objectness) * (1 - non_ignored_mask)
vessel_loss = vessel_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_is_vessel) * (1 - non_ignored_mask)
fishing_loss = fishing_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_is_fishing) * (1 - non_ignored_mask)
offset_loss = offset_loss * non_ignored_mask
size_loss = size_loss * non_ignored_mask

return (
    pos_objectness_loss,
    neg_objectness_loss,
    vessel_loss,
    fishing_loss,
    offset_loss,
    size_loss,
)
class YOLOLoss(nn.Module):
def init(self, num_classes: int, anchors: list):
super(YOLOLoss, self).init()
self.num_classes = num_classes
self.anchors = anchors
self.mse_loss = nn.MSELoss(reduction='sum') # Consider using 'mean' if preferred
self.bce_loss = nn.BCELoss(reduction='sum') # Consider using 'mean' if preferred

def calculate_iou(self, box1: torch.Tensor, box2: torch.Tensor) -> torch.Tensor:
    """Calculate Intersection Over Union (IoU) between two sets of boxes."""
    x1y1 = torch.max(box1[..., :2], box2[..., :2])
    x2y2 = torch.min(box1[..., 2:], box2[..., 2:])
    intersection = torch.clamp((x2y2 - x1y1), min=0).prod(-1)
    area_box1 = (box1[..., 2] - box1[..., 0]) * (box1[..., 3] - box1[..., 1])
    area_box2 = (box2[..., 2] - box2[..., 0]) * (box2[..., 3] - box2[..., 1])
    union = area_box1 + area_box2 - intersection
    return intersection / union

def forward(self, yolo_pred: torch.Tensor, yolo_targets: torch.Tensor) -> torch.Tensor:
    # Ensure predictions and targets have the same number of channels
    if yolo_pred.shape[1] != yolo_targets.shape[1]:
        raise ValueError(f"Channel mismatch: predictions have {yolo_pred.shape[1]} channels, targets have {yolo_targets.shape[1]} channels.")
    
    # Ensure spatial dimensions match
    if yolo_pred.shape[-2:] != yolo_targets.shape[-2:]:
        yolo_pred = F.interpolate(yolo_pred, size=yolo_targets.shape[-2:], mode='bilinear', align_corners=False)

    # Extract bounding boxes, confidence, and class predictions from YOLO output
    bbox_pred = yolo_pred[:, 0:4, :, :]  # [batch, 4, H, W]
    conf_pred = yolo_pred[:, 4:5, :, :]  # [batch, 1, H, W]
    class_pred = yolo_pred[:, 5:, :, :]  # [batch, num_classes, H, W]

    bbox_targets = yolo_targets[:, 0:4, :, :]  # [batch, 4, H, W]
    conf_targets = yolo_targets[:, 4:5, :, :]  # [batch, 1, H, W]
    class_targets = yolo_targets[:, 5:, :, :]  # [batch, num_classes, H, W]

    # Compute losses: bbox (MSE), confidence (BCE), and class (BCE or CrossEntropy)
    bbox_loss = self.mse_loss(bbox_pred, bbox_targets)
    conf_loss = self.bce_loss(conf_pred, conf_targets)
    class_loss = self.bce_loss(class_pred, class_targets)

    # Combine losses with appropriate weighting
    total_loss = bbox_loss + conf_loss + class_loss

    return total_loss
###original multilabel_circle_coder.py
import math
from collections import namedtuple
from typing import Tuple, Optional, Union
from torch.nn import functional as F
import numpy as np
import torch
from torch import Tensor
from .functional import *
all = [
"MultilabelCircleNetCoder",
"MultilabelCircleNetEncodeResult",
"MultilabelCircleNetDecodeResult",
]
from .functional import draw_circle
MultilabelCircleNetEncodeResult = namedtuple(
"MultilabelCircleNetEncodeResult",
["heatmap", "is_fishing", "is_vessel", "lengths", "offset"],
)
MultilabelCircleNetDecodeResult = namedtuple(
"MultilabelCircleNetDecodeResult",
["centers", "lengths", "is_fishing", "is_vessel", "scores"],
)
def cast_to_float_if_on_cpu(x: Tensor) -> Tensor:
if x.device == torch.device("cpu") and x.dtype == torch.float16:
x = x.float()
return x
class MultilabelCircleNetCoder:
output_stride: int
max_objects: int
heatmap_encoding: str
labels_encoding: str
ignore_value: Optional[int]
image_height: int
image_width: int
image_size: Tuple[int, int]
fixed_radius: int
labels_radius: int

def __init__(
    self,
    image_size: Tuple[int, int],
    output_stride: int,
    max_objects: int,
    heatmap_encoding: str,
    labels_encoding: str,
    ignore_value: Optional[int],
    fixed_radius: Optional[int] = None,
    labels_radius: Optional[int] = None,
    nms_kernel_size: int = 3,
    nms_method=centernet_heatmap_nms,
    ignore_low_confidence_detections: bool = False,
):
    image_size = image_size[:2]

    if heatmap_encoding not in {"umich", "msra"}:
        raise KeyError(heatmap_encoding)

    if labels_encoding not in {"point", "heatmap", "circle", "circle_ignore"}:
        raise KeyError(labels_encoding)

    self.heatmap_encoding = heatmap_encoding
    self.labels_encoding = labels_encoding
    self.output_stride = output_stride
    self.image_size = image_size
    self.image_height, self.image_width = image_size
    self.max_objects = max_objects
    self.ignore_value = ignore_value
    self.fixed_radius = fixed_radius
    self.labels_radius = labels_radius
    self.nms_kernel_size = nms_kernel_size
    self.ignore_low_confidence_detections = ignore_low_confidence_detections
    self.nms_method = nms_method

def __repr__(self):
    return (
        "MultilabelCircleNetCoder("
        + f"image_size={self.image_size}, "
        + f"stride={self.output_stride}, "
        + f"max_objects={self.max_objects}, "
        + f"labels_encoding={self.labels_encoding}, "
        + f"heatmap_encoding={self.heatmap_encoding}, "
        + f"fixed_radius={self.fixed_radius}, "
        + f"ignore_index={self.ignore_value}, "
        + f"ignore_low_confidence_detections={self.ignore_low_confidence_detections}"
        + ")"
    )

def box_coder_for_image_size(self, image_size):
    image_size = image_size[:2]
    if self.image_height == image_size[0] and self.image_width == image_size[1]:
        return self
    return MultilabelCircleNetCoder(
        image_size=image_size,
        output_stride=self.output_stride,
        max_objects=self.max_objects,
        heatmap_encoding=self.heatmap_encoding,
        labels_encoding=self.labels_encoding,
        ignore_value=self.ignore_value,
        fixed_radius=self.fixed_radius,
        labels_radius=self.labels_radius,
        nms_kernel_size=self.nms_kernel_size,
        nms_method=self.nms_method,
        ignore_low_confidence_detections=self.ignore_low_confidence_detections,
    )

def encode(
    self,
    centers: np.ndarray,
    confidences: np.ndarray,
    is_vessel: np.ndarray,
    is_fishing: np.ndarray,
    lengths: np.ndarray,
) -> MultilabelCircleNetEncodeResult:
    """
    :param centers [N,2] (x1,y1)
    :param lengths [N] (Ship length in pixels, See XView3DataModule.encode_lengths)
    """

    if (
        len(centers) != len(confidences)
        or len(centers) != len(is_vessel)
        or len(centers) != len(is_fishing)
        or len(centers) != len(lengths)
    ):
        raise RuntimeError(
            f"Got inconsistent length of centers ({centers.shape}), lengths ({lengths.shape}), is_vessel ({is_vessel.shape}) and is_fishing ({is_fishing.shape})"
        )

    # Sort objects by confidence
    low_conf_mask = confidences == "LOW"
    med_conf_mask = confidences == "MEDIUM"
    hgh_conf_mask = confidences == "HIGH"
    order = [low_conf_mask, med_conf_mask, hgh_conf_mask]

    def rearrange(x, masks, lengths):
        sorted_sub_arrays = []
        for m in masks:
            subsampled_lengths = lengths[m]
            subsampled_elements = x[m]
            order = np.argsort(-subsampled_lengths)
            sorted_sub_arrays.append(subsampled_elements[order])
        return np.concatenate(sorted_sub_arrays, axis=0)

    centers = rearrange(centers, order, lengths)
    confidences = rearrange(confidences, order, lengths)
    is_vessel = rearrange(is_vessel, order, lengths)
    is_fishing = rearrange(is_fishing, order, lengths)
    lengths = rearrange(lengths, order, lengths)

    # Reconstruct LOW mask after reordering
    low_conf_mask = confidences == "LOW"

    output_height = self.image_height // self.output_stride
    output_width = self.image_width // self.output_stride

    heatmap = np.zeros((1, output_height, output_width), dtype=np.float32)

    is_vessel_map = np.zeros((1, output_height, output_width), dtype=np.float32)
    is_fishing_map = np.zeros((1, output_height, output_width), dtype=np.float32)

    length_map = np.zeros((1, output_height, output_width), dtype=np.float32)
    offset_map = np.zeros((2, output_height, output_width), dtype=np.float32)

    num_objs = len(centers)

    centers_scaled = centers.astype(np.float32) / float(self.output_stride)
    lengths_scaled = lengths.astype(np.float32) / float(self.output_stride)

    log_lengths = self.length2log(lengths)  # Transform targets to log-space

    for i in range(num_objs):
        center = centers_scaled[i]
        center_int = center.astype(int)

        if self.fixed_radius:
            heatmap_radius = self.fixed_radius
        else:
            heatmap_radius = max(2, int(math.ceil(lengths_scaled[i] * 0.5)))

        if self.heatmap_encoding == "umich":
            if low_conf_mask[i] and self.ignore_low_confidence_detections:
                draw_circle(heatmap[0], center_int, heatmap_radius, self.ignore_value)
            else:
                draw_umich_gaussian(heatmap[0], center_int, heatmap_radius)
        elif self.heatmap_encoding == "point":
            if low_conf_mask[i] and self.ignore_low_confidence_detections:
                heatmap[0, center_int[1], center_int[0]] = self.ignore_value
            else:
                heatmap[0, center_int[1], center_int[0]] = 1
        else:
            raise KeyError(self.heatmap_encoding)

        if self.labels_encoding == "point":
            is_vessel_map[0, center_int[1], center_int[0]] = is_vessel[i]
            is_fishing_map[0, center_int[1], center_int[0]] = is_fishing[i]
            length_map[0, center_int[1], center_int[0]] = log_lengths[i]
        elif self.labels_encoding == "circle":
            if self.labels_radius:
                labels_radius = self.labels_radius
            else:
                labels_radius = heatmap_radius

            draw_circle(is_vessel_map[0], center_int, labels_radius, is_vessel[i])
            draw_circle(is_fishing_map[0], center_int, labels_radius, is_fishing[i])
            draw_circle(length_map[0], center_int, labels_radius, log_lengths[i])

        else:
            raise KeyError(self.labels_encoding)

        offset_map[0, center_int[1], center_int[0]] = center[0] - center_int[0]
        offset_map[1, center_int[1], center_int[0]] = center[1] - center_int[1]

    return MultilabelCircleNetEncodeResult(
        heatmap=heatmap,
        is_vessel=is_vessel_map,
        is_fishing=is_fishing_map,
        lengths=length_map,
        offset=offset_map,
    )

@torch.no_grad()
def decode(
    self,
    objectness_map: Tensor,
    is_vessel_map: Tensor,
    is_fishing_map: Tensor,
    length_map: Tensor,
    offset_map: Tensor,
    max_objects: Optional[int] = None,
    apply_activation: bool = False,
) -> MultilabelCircleNetDecodeResult:
    """
    Decode CenterNet predictions
    :param objectness_map: [B, 1, H, W]
    :param is_vessel_map: [B, 1, H, W]
    :param is_fishing_map: [B, 1, H, W]
    :param length_map: [B, 1, H, W]
    :param offset_map: [B, 2, H, W]
    :param max_objects: Maximum number of objects (K)
    :param apply_activation:
    :return: Tuple of 4 elements (centers, lengths, labels, scores)
        - [B, K, 2]
        - [B, K]
        - [B, K]
        - [B, K]
    """
    batch, _, height, width = objectness_map.size()

    objectness_map = cast_to_float_if_on_cpu(objectness_map)
    is_fishing_map = cast_to_float_if_on_cpu(is_fishing_map)
    is_vessel_map = cast_to_float_if_on_cpu(is_vessel_map)

    length_map = cast_to_float_if_on_cpu(length_map)

    if apply_activation:
        objectness_map = objectness_map.sigmoid()
        is_fishing_map = is_fishing_map.sigmoid()
        is_vessel_map = is_vessel_map.sigmoid()

    if max_objects is None:
        max_objects = self.max_objects

    # This is necessary if we use non-masked image (with NaNs)
    objectness_map = torch.masked_fill(objectness_map, ~torch.isfinite(objectness_map), 0)
    is_fishing_map = torch.masked_fill(is_fishing_map, ~torch.isfinite(is_fishing_map), 0)
    is_vessel_map = torch.masked_fill(is_vessel_map, ~torch.isfinite(is_vessel_map), 0)

    obj_map = self.nms_method(objectness_map, kernel=self.nms_kernel_size)

    # Limit K to prevent having K more W * H
    max_objects = min(max_objects, height * width)

    # Decode centers
    obj_scores, inds, _, ys, xs = centernet_topk(obj_map, top_k=max_objects)
    if offset_map is not None:
        offset_map = cast_to_float_if_on_cpu(offset_map)
        offset_map = centernet_tranpose_and_gather_feat(offset_map, inds)
        xs = xs.view(batch, max_objects, 1) + offset_map[:, :, 0:1]
        ys = ys.view(batch, max_objects, 1) + offset_map[:, :, 1:2]
    elif self.output_stride == 1:
        xs = xs.view(batch, max_objects, 1)
        ys = ys.view(batch, max_objects, 1)
    else:
        xs = xs.view(batch, max_objects, 1) + 0.5
        ys = ys.view(batch, max_objects, 1) + 0.5
    pred_centers = torch.cat([xs, ys], dim=2) * self.output_stride

    # Decode lengths
    pred_lengths = centernet_tranpose_and_gather_feat(length_map, inds)
    pred_lengths = pred_lengths[:, :, 0]  # Squeeze length

    # Decode is_vessel
    is_vessel, _ = centernet_tranpose_and_gather_feat(is_vessel_map, inds).max(dim=2)

    # Decode is_fishing
    is_fishing, _ = centernet_tranpose_and_gather_feat(is_fishing_map, inds).max(dim=2)

    return MultilabelCircleNetDecodeResult(
        centers=pred_centers,
        lengths=self.log2length(pred_lengths),  # Since length are in log-space, convert back to linear units
        is_fishing=is_fishing,
        is_vessel=is_vessel,
        scores=obj_scores,
    )

@classmethod
def log2length(cls, x: Union[Tensor]) -> Union[Tensor]:
    if torch.is_tensor(x):
        return torch.exp(F.relu(x)) - 1
    else:
        return np.exp(np.clip(x, a_min=0, a_max=None)) - 1

@classmethod
def length2log(cls, x: np.ndarray) -> np.ndarray:
    return np.log(x + 1)
my upgraded multilabel_circle_coder.py

import math
from collections import namedtuple
from typing import Tuple, Optional, Union, List, Dict, Any
from torch.nn import functional as F
import numpy as np
import torch
from torch import Tensor
from xview3.centernet.bboxer.functional import (
centernet_heatmap_nms,
centernet_tight_heatmap_nms,
centernet_topk,
centernet_tranpose_and_gather_feat,
draw_umich_gaussian,
)
from .functional import draw_circle
all = [
"MultilabelCircleNetCoder",
"MultilabelCircleNetEncodeResult",
"MultilabelCircleNetDecodeResult",
]
MultilabelCircleNetEncodeResult = namedtuple(
"MultilabelCircleNetEncodeResult",
["heatmap", "is_fishing", "is_vessel", "lengths", "offset", "yolo_targets"],
)
MultilabelCircleNetDecodeResult = namedtuple(
"MultilabelCircleNetDecodeResult",
["centers", "lengths", "is_fishing", "is_vessel", "scores"],
)
Define available NMS methods
NMS_METHODS = {
'centernet_heatmap_nms': centernet_heatmap_nms,
'centernet_tight_heatmap_nms': centernet_tight_heatmap_nms,
}
def cast_to_float_if_on_cpu(x: Tensor) -> Tensor:

if x.device == torch.device("cpu") and x.dtype == torch.float16:
    x = x.float()
return x
class MultilabelCircleNetCoder:
def init(
self,
image_size: Tuple[int, int],
output_stride: int,
max_objects: int,
heatmap_encoding: str,
labels_encoding: str,
ignore_value: Optional[int],
fixed_radius: Optional[int] = None,
labels_radius: Optional[int] = None,
nms_kernel_size: int = 3,
nms_method=centernet_heatmap_nms,
ignore_low_confidence_detections: bool = False,
):

    image_size = image_size[:2]

    if heatmap_encoding not in {"umich", "msra"}:
        raise KeyError(heatmap_encoding)

    if labels_encoding not in {"point", "heatmap", "circle", "circle_ignore"}:
        raise KeyError(labels_encoding)

    self.image_size = image_size
    self.output_stride = output_stride
    self.max_objects = max_objects
    self.heatmap_encoding = heatmap_encoding
    self.labels_encoding = labels_encoding
    self.ignore_value = ignore_value
    self.fixed_radius = fixed_radius
    self.labels_radius = labels_radius
    self.nms_kernel_size = nms_kernel_size
    self.nms_method = nms_method
    self.ignore_low_confidence_detections = ignore_low_confidence_detections
    self.image_height, self.image_width = image_size
    self.grid_height = self.image_height // self.output_stride
    self.grid_width = self.image_width // self.output_stride

    # Define mappings for categorical fields
    self.confidence_mapping = {
        'NONE': 0.0,
        'LOW': 0.33,
        'MEDIUM': 0.66,
        'HIGH': 1.0,
    }
    self.vessel_mapping = {
        'UNKNOWN': 0,
        'TYPE1': 1,
        'TYPE2': 2,
        # Add more vessel types as necessary
    }
    self.fishing_mapping = {
        'UNKNOWN': 0,
        'ACTIVE': 1,
        'INACTIVE': 2,
        # Add more fishing statuses as necessary
    }

def to_dict(self) -> Dict[str, Any]:
    
    return {
        "type": "MultilabelCircleNetCoder",
        "image_size": self.image_size,
        "max_objects": self.max_objects,
        "heatmap_encoding": self.heatmap_encoding,
        "labels_encoding": self.labels_encoding,
        "ignore_value": self.ignore_value,
        "output_stride": self.output_stride,
        "fixed_radius": self.fixed_radius,
        "labels_radius": self.labels_radius,
        "nms_kernel_size": self.nms_kernel_size,
        "nms_method": self.nms_method,
        "ignore_low_confidence_detections": self.ignore_low_confidence_detections
    }

def __repr__(self) -> str:
    return (
        "MultilabelCircleNetCoder("
        + f"image_size={self.image_size}, "
        + f"stride={self.output_stride}, "
        + f"max_objects={self.max_objects}, "
        + f"labels_encoding={self.labels_encoding}, "
        + f"heatmap_encoding={self.heatmap_encoding}, "
        + f"fixed_radius={self.fixed_radius}, "
        + f"ignore_index={self.ignore_value}, "
        + f"ignore_low_confidence_detections={self.ignore_low_confidence_detections}"
        + ")"
    )

def create_empty_targets(self) -> MultilabelCircleNetEncodeResult:
    # Initialize empty or zero-filled arrays with the correct dimensions
    heatmap = np.zeros((1, self.grid_height, self.grid_width), dtype=np.float32)
    is_vessel = np.zeros_like(heatmap)
    is_fishing = np.zeros_like(heatmap)
    lengths = np.zeros_like(heatmap)
    offset = np.zeros((2, self.grid_height, self.grid_width), dtype=np.float32)
    yolo_targets = np.zeros((self.grid_height, self.grid_width, 6), dtype=np.float32)  # Adjust as needed

    return MultilabelCircleNetEncodeResult(
        heatmap=heatmap,
        is_vessel=is_vessel,
        is_fishing=is_fishing,
        lengths=lengths,
        offset=offset,
        yolo_targets=yolo_targets
    )



def create_box_coder_for_image_size(self, image_size: Optional[Tuple[int, int]] = None) -> 'MultilabelCircleNetCoder':
    
    if image_size is None:
        image_size = self.image_size
    return MultilabelCircleNetCoder(
        image_size=image_size,
        output_stride=self.output_stride,
        max_objects=self.max_objects,
        heatmap_encoding=self.heatmap_encoding,
        labels_encoding=self.labels_encoding,
        ignore_value=self.ignore_value,
        fixed_radius=self.fixed_radius,
        labels_radius=self.labels_radius,
        nms_kernel_size=self.nms_kernel_size,
        nms_method=self.nms_method,
        ignore_low_confidence_detections=self.ignore_low_confidence_detections,
    )

def encode(
    self,
    centers: np.ndarray,
    confidences: List[str],
    is_vessel: List[str],
    is_fishing: List[str],
    lengths: np.ndarray,
) -> MultilabelCircleNetEncodeResult:
   
    if centers.size == 0:
        return self.create_empty_targets()

    # Validate input lengths
    if (
        len(centers) != len(confidences)
        or len(centers) != len(is_vessel)
        or len(centers) != len(is_fishing)
        or len(centers) != len(lengths)
    ):
        raise RuntimeError(
            f"Got inconsistent length of centers ({centers.shape}), lengths ({lengths.shape}), is_vessel ({len(is_vessel)}) and is_fishing ({len(is_fishing)})"
        )

        # Convert lists to NumPy arrays
    confidences = np.array(confidences)
    is_vessel = np.array(is_vessel)
    is_fishing = np.array(is_fishing)

    # Convert confidences to a flat list of strings
    if confidences.ndim > 1:
        # Flatten the array if necessary
        confidences_list = confidences.flatten().tolist()
    else:
        confidences_list = confidences.tolist()

    # Ensure all elements are strings
    confidences_list = [str(elem) for elem in confidences_list]

    expected_confidences = {"LOW", "MEDIUM", "HIGH", "UNKNOWN"}

    # Validate categorical fields
    if not set(confidences_list).issubset(expected_confidences):
        unexpected = set(confidences_list) - expected_confidences
        raise ValueError(f"Confidences contain unexpected values: {unexpected}")

    # Sort objects by confidence: LOW < MEDIUM < HIGH < UNKNOWN
    low_conf_mask = confidences == "LOW"
    med_conf_mask = confidences == "MEDIUM"
    hgh_conf_mask = confidences == "HIGH"
    unknown_conf_mask = confidences == "UNKNOWN"

    order = [low_conf_mask, med_conf_mask, hgh_conf_mask, unknown_conf_mask]

    def rearrange(x, masks, lengths):
        sorted_sub_arrays = []
        for m in masks:
            if not isinstance(m, np.ndarray) or m.dtype != bool:
                raise TypeError("Mask must be a NumPy boolean array")
            subsampled_lengths = lengths[m]
            subsampled_elements = x[m]
            order = np.argsort(-subsampled_lengths)  # Sort descending by length
            sorted_sub_arrays.append(subsampled_elements[order])
        if sorted_sub_arrays:
            return np.concatenate(sorted_sub_arrays, axis=0)
        else:
            return np.array([]).reshape(0, x.shape[1]) if x.ndim == 2 else np.array([]).reshape(0)

    centers_sorted = rearrange(centers, order, lengths)
    confidences_sorted = rearrange(confidences, order, lengths)
    is_vessel_sorted = rearrange(is_vessel, order, lengths)
    is_fishing_sorted = rearrange(is_fishing, order, lengths)
    lengths_sorted = rearrange(lengths, order, lengths)

    # Create heatmap
    heatmap = np.zeros((1, self.grid_height, self.grid_width), dtype=np.float32)

    # Initialize maps
    is_vessel_map = np.zeros((1, self.grid_height, self.grid_width), dtype=np.float32)
    is_fishing_map = np.zeros((1, self.grid_height, self.grid_width), dtype=np.float32)
    length_map = np.zeros((1, self.grid_height, self.grid_width), dtype=np.float32)
    offset_map = np.zeros((2, self.grid_height, self.grid_width), dtype=np.float32)

    num_objs = len(centers_sorted)

    centers_scaled = centers_sorted.astype(np.float32) / float(self.output_stride)
    lengths_scaled = lengths_sorted.astype(np.float32) / float(self.output_stride)

    log_lengths = self.length2log(lengths_sorted)  # Transform targets to log-space

    for i in range(num_objs):
        center = centers_scaled[i]
        center_int = center.astype(int)

        # Determine radius
        if self.fixed_radius:
            heatmap_radius = self.fixed_radius
        else:
            heatmap_radius = max(2, int(math.ceil(lengths_scaled[i] * 0.5)))

        # Generate heatmap
        if self.heatmap_encoding == "umich":
            if self.labels_encoding == "circle_ignore" and self.ignore_low_confidence_detections:
                draw_circle(heatmap[0], center_int, heatmap_radius, self.ignore_value)
            else:
                draw_umich_gaussian(heatmap[0], center_int, heatmap_radius)
        elif self.heatmap_encoding == "point":
            if self.labels_encoding == "circle_ignore" and self.ignore_low_confidence_detections:
                heatmap[0, center_int[1], center_int[0]] = self.ignore_value
            else:
                heatmap[0, center_int[1], center_int[0]] = 1.0
        else:
            raise KeyError(f"Unknown heatmap_encoding: {self.heatmap_encoding}")

        # Encode vessel and fishing maps
        if self.labels_encoding in {"circle", "circle_ignore"}:
            labels_radius = self.labels_radius if self.labels_radius else heatmap_radius
            draw_circle(is_vessel_map[0], center_int, labels_radius, self.vessel_mapping.get(is_vessel_sorted[i], self.ignore_value))
            draw_circle(is_fishing_map[0], center_int, labels_radius, self.fishing_mapping.get(is_fishing_sorted[i], self.ignore_value))
        elif self.labels_encoding == "point":
            # Handle point encoding if necessary
            raise NotImplementedError(f"Labels encoding 'point' not implemented")
        else:
            raise KeyError(f"Unknown labels_encoding: {self.labels_encoding}")

        # Encode lengths
        length_map[0, center_int[1], center_int[0]] = log_lengths[i]

        # Encode offset
        offset_map[0, center_int[1], center_int[0]] = center[0] - center_int[0]
        offset_map[1, center_int[1], center_int[0]] = center[1] - center_int[1]

    # Generate YOLO targets
    yolo_targets = self.generate_yolo_targets(centers_sorted, confidences_sorted, is_vessel_sorted, is_fishing_sorted, lengths_sorted)
    
    return MultilabelCircleNetEncodeResult(
        heatmap=heatmap,
        is_fishing=is_fishing_map,            
        is_vessel=is_vessel_map,
        lengths=length_map,
        offset=offset_map,
        yolo_targets=yolo_targets,  # Include yolo_targets in the result
    )

def generate_yolo_targets(
    self,
    centers: np.ndarray,
    confidences: np.ndarray,
    is_vessel: np.ndarray,
    is_fishing: np.ndarray,
    lengths: np.ndarray
) -> np.ndarray:
    
    num_objects = len(centers)
    yolo_targets = np.zeros((self.grid_height, self.grid_width, 6), dtype=np.float32)

    for i in range(num_objects):
        center = centers[i]
        confidence = self.confidence_mapping.get(confidences[i], self.ignore_value)  # Convert confidence to numerical value
        vessel_type = self.vessel_mapping.get(is_vessel[i], self.ignore_value)  # Ensure correct mapping
        fishing_status = self.fishing_mapping.get(is_fishing[i], self.ignore_value)  # Ensure correct mapping
        length = lengths[i]

        # Convert center coordinates to grid cell format (relative to the grid)
        grid_x = int(center[0] / self.output_stride)
        grid_y = int(center[1] / self.output_stride)

        # Ensure grid indices are within bounds
        grid_x = min(grid_x, self.grid_width - 1)
        grid_y = min(grid_y, self.grid_height - 1)

        # Assign to the correct grid cell
        yolo_targets[grid_y, grid_x, 0:4] = np.array([center[0], center[1], length, length])
        yolo_targets[grid_y, grid_x, 4] = confidence
        yolo_targets[grid_y, grid_x, 5] = vessel_type

    return yolo_targets

@torch.no_grad()
def decode(
    self,
    objectness_map: Tensor,
    is_vessel_map: Tensor,
    is_fishing_map: Tensor,
    length_map: Tensor,
    offset_map: Tensor,
    max_objects: Optional[int] = None,
    apply_activation: bool = False,
) -> MultilabelCircleNetDecodeResult:
   
    batch, _, height, width = objectness_map.size()

    objectness_map = cast_to_float_if_on_cpu(objectness_map)
    is_fishing_map = cast_to_float_if_on_cpu(is_fishing_map)
    is_vessel_map = cast_to_float_if_on_cpu(is_vessel_map)

    length_map = cast_to_float_if_on_cpu(length_map)

    if apply_activation:
        objectness_map = objectness_map.sigmoid()
        is_fishing_map = is_fishing_map.sigmoid()
        is_vessel_map = is_vessel_map.sigmoid()

    if max_objects is None:
        max_objects = self.max_objects

    # Handle non-masked images (with NaNs)
    objectness_map = torch.masked_fill(objectness_map, ~torch.isfinite(objectness_map), 0)
    is_fishing_map = torch.masked_fill(is_fishing_map, ~torch.isfinite(is_fishing_map), 0)
    is_vessel_map = torch.masked_fill(is_vessel_map, ~torch.isfinite(is_vessel_map), 0)

    obj_map = self.nms_method(objectness_map, kernel=self.nms_kernel_size)

    # Limit K to prevent having K more than W * H
    max_objects = min(max_objects, height * width)

    # Decode centers
    obj_scores, inds, _, ys, xs = centernet_topk(obj_map, top_k=max_objects)
    if offset_map is not None:
        offset_map = cast_to_float_if_on_cpu(offset_map)
        offset_map = centernet_tranpose_and_gather_feat(offset_map, inds)
        xs = xs.view(batch, max_objects, 1) + offset_map[:, :, 0:1]
        ys = ys.view(batch, max_objects, 1) + offset_map[:, :, 1:2]
    elif self.output_stride == 1:
        xs = xs.view(batch, max_objects, 1)
        ys = ys.view(batch, max_objects, 1)
    else:
        xs = xs.view(batch, max_objects, 1) + 0.5
        ys = ys.view(batch, max_objects, 1) + 0.5
    pred_centers = torch.cat([xs, ys], dim=2) * self.output_stride

    # Decode lengths
    pred_lengths = centernet_tranpose_and_gather_feat(length_map, inds)
    pred_lengths = pred_lengths[:, :, 0]  # Squeeze length

    # Decode is_vessel
    is_vessel = torch.sigmoid(centernet_tranpose_and_gather_feat(is_vessel_map, inds))  # Apply sigmoid for binary classification
    is_vessel = is_vessel.max(dim=2)[0]

    # Decode is_fishing
    is_fishing = torch.sigmoid(centernet_tranpose_and_gather_feat(is_fishing_map, inds))  # Apply sigmoid for binary classification
    is_fishing = is_fishing.max(dim=2)[0]

    return MultilabelCircleNetDecodeResult(
        centers=pred_centers,
        lengths=self.log2length(pred_lengths),  # Since lengths are in log-space, convert back to linear units
        is_fishing=is_fishing,
        is_vessel=is_vessel,
        scores=obj_scores,
    )

@classmethod
def from_config(cls, config: Dict[str, Any]) -> 'MultilabelCircleNetCoder':
    
    nms_method_name = config['nms_method']
    nms_method = NMS_METHODS.get(nms_method_name)
    
    if nms_method is None:
        raise ValueError(f"Unknown nms_method: {nms_method_name}")
    
    return cls(
        image_size=config['image_size'],
        output_stride=config['output_stride'],
        max_objects=config['max_objects'],
        heatmap_encoding=config['heatmap_encoding'],
        labels_encoding=config['labels_encoding'],
        ignore_value=config['ignore_value'],
        fixed_radius=config.get('fixed_radius'),
        labels_radius=config.get('labels_radius'),
        nms_kernel_size=config.get('nms_kernel_size', 3),
        nms_method=nms_method,
        ignore_low_confidence_detections=config.get('ignore_low_confidence_detections', False),
    )

@classmethod
def log2length(cls, x: Union[Tensor, np.ndarray]) -> Union[Tensor, np.ndarray]:
    
    if torch.is_tensor(x):
        return torch.exp(F.relu(x)) - 1
    else:
        return np.exp(np.clip(x, a_min=0, a_max=None)) - 1

@classmethod
def length2log(cls, x: np.ndarray) -> np.ndarray:
    
    return np.log(x + 1)

It looks like you're attempting to integrate YOLO-style target generation and loss calculation into your existing CenterNet-based multi-label detection code. This is a complex task and can introduce critical issues if not carefully managed.
After reviewing your upgraded versions of multilabel_trimmed.loss.py and multilabel_circle_coder.py, I noticed several areas where errors are likely occurring. I'll walk through these areas and provide guidance on how to fix them.
1. Issues in multilabel_trimmed.loss.py

a. Incorrect Dictionary Access
In your forward method, you're accessing kwargs incorrectly:

gt_yolo: Tensor = kwargs(CENTERNET_TARGET_YOLO)
Error: You're using parentheses instead of square brackets to access dictionary items. This will result in a TypeError.
Fix:

gt_yolo: Tensor = kwargs.get(CENTERNET_TARGET_YOLO)
if gt_yolo is None:
    raise ValueError(f"Missing key '{CENTERNET_TARGET_YOLO}' in input dictionary")
Similarly, for pred_yolo, you should use:

pred_yolo: Tensor = kwargs.get(CENTERNET_OUTPUT_YOLO, torch.zeros_like(pred_objectness))
b. YOLO Loss Integration
You've added a YOLOLoss class and instantiated it within your MultilabelCenterTrimmedLoss. However, there are issues:
Unused anchors and num_classes: While you pass num_classes and anchors to your YOLOLoss, you're not utilizing them within the class. In YOLO, anchors are crucial for predicting bounding boxes, and num_classes is essential for handling multiple classes.
Fix:
You need to adjust your YOLOLoss to use anchors and num_classes appropriately. For example, you might need to:
Use anchors to adjust the predicted bounding box coordinates.
Use num_classes to set up the final layer dimensions and compute the classification loss correctly.
c. Data Shape and Dimension Mismatches
In your YOLOLoss.forward method, you're assuming that yolo_pred and yolo_targets have shapes [batch_size, channels, height, width], but in your encode method in multilabel_circle_coder.py, yolo_targets are generated as numpy arrays of shape [grid_height, grid_width, 6].
Fix:
Ensure that both yolo_pred and yolo_targets have matching shapes and dimensions. You might need to transpose or reshape your yolo_targets in your dataset class to match the model output shape.
Here's how you can adjust the shape:

# In your MultilabelCircleNetCoder.encode method
yolo_targets = yolo_targets.transpose(2, 0, 1)  # Now shape is [6, grid_height, grid_width]
Then, in your dataset's __getitem__ method, you should ensure that yolo_targets have a batch dimension:

yolo_targets = torch.from_numpy(yolo_targets).float()  # [6, H, W]
If required, unsqueeze to add a batch dimension:

yolo_targets = yolo_targets.unsqueeze(0)  # [1, 6, H, W]
d. Loss Function Adjustments
Your YOLOLoss class needs to compute the loss components appropriately. Currently, it uses MSELoss and BCELoss without considering the proper YOLO loss components.
Fix:
In YOLO, the loss is typically a combination of:
Localization Loss: Often using MSELoss for bounding box coordinates.
Confidence Loss: Using BCELoss for object confidence score.
Classification Loss: Depending on the number of classes, you might use CrossEntropyLoss (for multi-class) or BCELoss (for multi-label).
Example Adjustments:

class YOLOLoss(nn.Module):
    def __init__(self, num_classes: int, anchors: list):
        super(YOLOLoss, self).__init__()
        self.num_classes = num_classes
        self.anchors = anchors
        self.mse_loss = nn.MSELoss(reduction='sum')  # For location loss
        self.bce_loss = nn.BCELoss(reduction='sum')  # For confidence loss
        self.class_loss_fn = nn.CrossEntropyLoss(reduction='sum')  # For classification loss

    def forward(self, yolo_pred: torch.Tensor, yolo_targets: torch.Tensor) -> torch.Tensor:
        # Ensure predictions and targets have the same shape
        if yolo_pred.shape != yolo_targets.shape:
            raise ValueError(f"Shape mismatch: predictions have shape {yolo_pred.shape}, targets have shape {yolo_targets.shape}.")

        # Extract components
        bbox_pred = yolo_pred[:, :4, :, :]        # [batch, 4, H, W]
        conf_pred = yolo_pred[:, 4, :, :]         # [batch, H, W]
        class_pred = yolo_pred[:, 5:, :, :]       # [batch, num_classes, H, W]

        bbox_targets = yolo_targets[:, :4, :, :]  # [batch, 4, H, W]
        conf_targets = yolo_targets[:, 4, :, :]   # [batch, H, W]
        class_targets = yolo_targets[:, 5:, :, :] # [batch, num_classes, H, W] or [batch, H, W] if using class indices

        # Compute losses
        bbox_loss = self.mse_loss(bbox_pred, bbox_targets)
        conf_loss = self.bce_loss(conf_pred, conf_targets)

        # Adjust class loss computation based on whether it's multi-class or multi-label
        # For multi-class classification (class indices)
        # class_pred: [batch, num_classes, H, W]
        # class_targets: [batch, H, W], each element is a class index
        class_pred = class_pred.permute(0, 2, 3, 1).contiguous()  # [batch, H, W, num_classes]
        class_pred = class_pred.view(-1, self.num_classes)        # [batch*H*W, num_classes]
        class_targets = class_targets.view(-1)                    # [batch*H*W]
        class_loss = self.class_loss_fn(class_pred, class_targets)

        # Combine losses
        total_loss = bbox_loss + conf_loss + class_loss
        return total_loss
2. Issues in multilabel_circle_coder.py

a. Generating YOLO Targets
In your generate_yolo_targets method, you're creating targets of shape [grid_height, grid_width, 6].
Fix:
Ensure that your yolo_targets are consistent with the rest of your targets in terms of dimensions and data types. Here's how you can adjust:

# After generating yolo_targets, transpose dimensions to match [channels, H, W]
yolo_targets = yolo_targets.transpose(2, 0, 1)  # Shape: [6, grid_height, grid_width]
If required, ensure that yolo_targets is converted to a torch.Tensor and has the batch dimension:

yolo_targets = torch.from_numpy(yolo_targets).float()  # Convert to torch.Tensor
yolo_targets = yolo_targets.unsqueeze(0)               # Add batch dimension if necessary
b. Data Type Consistency
Be consistent with data types. Make sure that all your numerical data are in float32 and categorical data are appropriately encoded as integers or one-hot vectors.
c. Coordinate Scaling
In your generate_yolo_targets method, when you're computing grid_x and grid_y, you need to ensure that you're scaling the centers correctly, considering the output_stride.
Example Fix:

grid_x = int(center[0] // self.output_stride)
grid_y = int(center[1] // self.output_stride)
3. Integration with Dataset and Collation

Ensure that your dataset class and centernet_collate function correctly handle the new CENTERNET_TARGET_YOLO key.
In your dataset class (e.g., MultilabelCircleNetRandomCropDataset):
Include yolo_targets in the sample dictionary returned by __getitem__.
In your centernet_collate function:
Add CENTERNET_TARGET_YOLO to the skip_keys list if you want to process it separately.
4. Testing and Validation

a. Unit Testing
Before integrating everything, test each component individually:
Test MultilabelCircleNetCoder.encode to see if it produces correctly shaped outputs.
Validate that yolo_targets are correctly generated and have the expected dimensions.
Test YOLOLoss with dummy data to ensure it computes the loss without errors.
b. Small Batch Testing
Run your data loader and model forward pass with a small batch (e.g., batch size of 1 or 2) to identify any runtime errors.
5. Considerations for Integrating YOLO into CenterNet

Integrating YOLO into a CenterNet-based architecture is non-trivial due to fundamental differences in the detection paradigms:
CenterNet focuses on keypoint detection and uses heatmaps to identify object centers.
YOLO predicts bounding boxes and class probabilities directly over a grid.
Mixing these methods can lead to conflicts in the data processing pipeline, model architecture, and loss functions.
Suggestion:
Carefully evaluate whether integrating YOLO components into your CenterNet model is the best approach.
If your goal is to leverage YOLO's capabilities, consider using a YOLO-based model and adjust it for your multi-label detection task.
Alternatively, stick with the CenterNet architecture and enhance it according to your specific requirements.
6. Next Steps

Fix the identified errors in your code as described above.
Ensure data shapes and types are consistent throughout your pipeline.
Test each component individually to isolate and fix issues.
Reconsider the integration strategy between YOLO and CenterNet, potentially choosing one architecture to align with your task.
If you continue to encounter issues, please provide:
Specific error messages and stack traces.
Details about your model's expected input and output shapes.
Any relevant snippets of code where errors occur.
This information will help in providing more precise guidance.
I hope this helps you identify and fix the critical issues in your code. Integrating complex architectures requires careful adjustments, but with systematic troubleshooting, you can resolve these problems.

perfect I've applied your suggestions and fixes precisely put can you review the upgraded version of multilabel_circle_coder.py and ensure the compatibility after integrating yolo with multilabel_dataset.py
here's my fixed
multilabel_circle_coder.py
import math
from collections import namedtuple
from typing import Tuple, Optional, Union, List, Dict, Any
from torch.nn import functional as F
import numpy as np
import torch
from torch import Tensor
from xview3.centernet.bboxer.functional import (
centernet_heatmap_nms,
centernet_tight_heatmap_nms,
centernet_topk,
centernet_tranpose_and_gather_feat,
draw_umich_gaussian,
)
from .functional import draw_circle
all = [
"MultilabelCircleNetCoder",
"MultilabelCircleNetEncodeResult",
"MultilabelCircleNetDecodeResult",
]
MultilabelCircleNetEncodeResult = namedtuple(
"MultilabelCircleNetEncodeResult",
["heatmap", "is_fishing", "is_vessel", "lengths", "offset", "yolo_targets"],
)
MultilabelCircleNetDecodeResult = namedtuple(
"MultilabelCircleNetDecodeResult",
["centers", "lengths", "is_fishing", "is_vessel", "scores"],
)
Define available NMS methods
NMS_METHODS = {
'centernet_heatmap_nms': centernet_heatmap_nms,
'centernet_tight_heatmap_nms': centernet_tight_heatmap_nms,
}
def cast_to_float_if_on_cpu(x: Tensor) -> Tensor:

if x.device == torch.device("cpu") and x.dtype == torch.float16:
    x = x.float()
return x
class MultilabelCircleNetCoder:
def init(
self,
image_size: Tuple[int, int],
output_stride: int,
max_objects: int,
heatmap_encoding: str,
labels_encoding: str,
ignore_value: Optional[int],
fixed_radius: Optional[int] = None,
labels_radius: Optional[int] = None,
nms_kernel_size: int = 3,
nms_method=centernet_heatmap_nms,
ignore_low_confidence_detections: bool = False,
):

    image_size = image_size[:2]

    if heatmap_encoding not in {"umich", "msra"}:
        raise KeyError(heatmap_encoding)

    if labels_encoding not in {"point", "heatmap", "circle", "circle_ignore"}:
        raise KeyError(labels_encoding)

    self.image_size = image_size
    self.output_stride = output_stride
    self.max_objects = max_objects
    self.heatmap_encoding = heatmap_encoding
    self.labels_encoding = labels_encoding
    self.ignore_value = ignore_value
    self.fixed_radius = fixed_radius
    self.labels_radius = labels_radius
    self.nms_kernel_size = nms_kernel_size
    self.nms_method = nms_method
    self.ignore_low_confidence_detections = ignore_low_confidence_detections
    self.image_height, self.image_width = image_size
    self.grid_height = self.image_height // self.output_stride
    self.grid_width = self.image_width // self.output_stride

    # Define mappings for categorical fields
    self.confidence_mapping = {
        'NONE': 0.0,
        'LOW': 0.33,
        'MEDIUM': 0.66,
        'HIGH': 1.0,
    }
    self.vessel_mapping = {
        'UNKNOWN': 0,
        'TYPE1': 1,
        'TYPE2': 2,
        # Add more vessel types as necessary
    }
    self.fishing_mapping = {
        'UNKNOWN': 0,
        'ACTIVE': 1,
        'INACTIVE': 2,
        # Add more fishing statuses as necessary
    }

def to_dict(self) -> Dict[str, Any]:
    
    return {
        "type": "MultilabelCircleNetCoder",
        "image_size": self.image_size,
        "max_objects": self.max_objects,
        "heatmap_encoding": self.heatmap_encoding,
        "labels_encoding": self.labels_encoding,
        "ignore_value": self.ignore_value,
        "output_stride": self.output_stride,
        "fixed_radius": self.fixed_radius,
        "labels_radius": self.labels_radius,
        "nms_kernel_size": self.nms_kernel_size,
        "nms_method": self.nms_method,
        "ignore_low_confidence_detections": self.ignore_low_confidence_detections
    }

def __repr__(self) -> str:
    return (
        "MultilabelCircleNetCoder("
        + f"image_size={self.image_size}, "
        + f"stride={self.output_stride}, "
        + f"max_objects={self.max_objects}, "
        + f"labels_encoding={self.labels_encoding}, "
        + f"heatmap_encoding={self.heatmap_encoding}, "
        + f"fixed_radius={self.fixed_radius}, "
        + f"ignore_index={self.ignore_value}, "
        + f"ignore_low_confidence_detections={self.ignore_low_confidence_detections}"
        + ")"
    )

def create_empty_targets(self) -> MultilabelCircleNetEncodeResult:
    # Initialize empty or zero-filled arrays with the correct dimensions
    heatmap = np.zeros((1, self.grid_height, self.grid_width), dtype=np.float32)
    is_vessel = np.zeros_like(heatmap)
    is_fishing = np.zeros_like(heatmap)
    lengths = np.zeros_like(heatmap)
    offset = np.zeros((2, self.grid_height, self.grid_width), dtype=np.float32)
    yolo_targets = np.zeros((self.grid_height, self.grid_width, 6), dtype=np.float32)  # Adjust as needed

    return MultilabelCircleNetEncodeResult(
        heatmap=heatmap,
        is_vessel=is_vessel,
        is_fishing=is_fishing,
        lengths=lengths,
        offset=offset,
        yolo_targets=yolo_targets
    )



def create_box_coder_for_image_size(self, image_size: Optional[Tuple[int, int]] = None) -> 'MultilabelCircleNetCoder':
    
    if image_size is None:
        image_size = self.image_size
    return MultilabelCircleNetCoder(
        image_size=image_size,
        output_stride=self.output_stride,
        max_objects=self.max_objects,
        heatmap_encoding=self.heatmap_encoding,
        labels_encoding=self.labels_encoding,
        ignore_value=self.ignore_value,
        fixed_radius=self.fixed_radius,
        labels_radius=self.labels_radius,
        nms_kernel_size=self.nms_kernel_size,
        nms_method=self.nms_method,
        ignore_low_confidence_detections=self.ignore_low_confidence_detections,
    )

def encode(
    self,
    centers: np.ndarray,
    confidences: List[str],
    is_vessel: List[str],
    is_fishing: List[str],
    lengths: np.ndarray,
) -> MultilabelCircleNetEncodeResult:
   
    if centers.size == 0:
        return self.create_empty_targets()

    # Validate input lengths
    if (
        len(centers) != len(confidences)
        or len(centers) != len(is_vessel)
        or len(centers) != len(is_fishing)
        or len(centers) != len(lengths)
    ):
        raise RuntimeError(
            f"Got inconsistent length of centers ({centers.shape}), lengths ({lengths.shape}), is_vessel ({len(is_vessel)}) and is_fishing ({len(is_fishing)})"
        )

        # Convert lists to NumPy arrays
    confidences = np.array(confidences)
    is_vessel = np.array(is_vessel)
    is_fishing = np.array(is_fishing)

    # Convert confidences to a flat list of strings
    if confidences.ndim > 1:
        # Flatten the array if necessary
        confidences_list = confidences.flatten().tolist()
    else:
        confidences_list = confidences.tolist()

    # Ensure all elements are strings
    confidences_list = [str(elem) for elem in confidences_list]

    expected_confidences = {"LOW", "MEDIUM", "HIGH", "UNKNOWN"}

    # Validate categorical fields
    if not set(confidences_list).issubset(expected_confidences):
        unexpected = set(confidences_list) - expected_confidences
        raise ValueError(f"Confidences contain unexpected values: {unexpected}")

    # Sort objects by confidence: LOW < MEDIUM < HIGH < UNKNOWN
    low_conf_mask = confidences == "LOW"
    med_conf_mask = confidences == "MEDIUM"
    hgh_conf_mask = confidences == "HIGH"
    unknown_conf_mask = confidences == "UNKNOWN"

    order = [low_conf_mask, med_conf_mask, hgh_conf_mask, unknown_conf_mask]

    def rearrange(x, masks, lengths):
        sorted_sub_arrays = []
        for m in masks:
            if not isinstance(m, np.ndarray) or m.dtype != bool:
                raise TypeError("Mask must be a NumPy boolean array")
            subsampled_lengths = lengths[m]
            subsampled_elements = x[m]
            order = np.argsort(-subsampled_lengths)  # Sort descending by length
            sorted_sub_arrays.append(subsampled_elements[order])
        if sorted_sub_arrays:
            return np.concatenate(sorted_sub_arrays, axis=0)
        else:
            return np.array([]).reshape(0, x.shape[1]) if x.ndim == 2 else np.array([]).reshape(0)

    centers_sorted = rearrange(centers, order, lengths)
    confidences_sorted = rearrange(confidences, order, lengths)
    is_vessel_sorted = rearrange(is_vessel, order, lengths)
    is_fishing_sorted = rearrange(is_fishing, order, lengths)
    lengths_sorted = rearrange(lengths, order, lengths)

    # Create heatmap
    heatmap = np.zeros((1, self.grid_height, self.grid_width), dtype=np.float32)

    # Initialize maps
    is_vessel_map = np.zeros((1, self.grid_height, self.grid_width), dtype=np.float32)
    is_fishing_map = np.zeros((1, self.grid_height, self.grid_width), dtype=np.float32)
    length_map = np.zeros((1, self.grid_height, self.grid_width), dtype=np.float32)
    offset_map = np.zeros((2, self.grid_height, self.grid_width), dtype=np.float32)

    num_objs = len(centers_sorted)

    centers_scaled = centers_sorted.astype(np.float32) / float(self.output_stride)
    lengths_scaled = lengths_sorted.astype(np.float32) / float(self.output_stride)

    log_lengths = self.length2log(lengths_sorted)  # Transform targets to log-space

    for i in range(num_objs):
        center = centers_scaled[i]
        center_int = center.astype(int)

        # Determine radius
        if self.fixed_radius:
            heatmap_radius = self.fixed_radius
        else:
            heatmap_radius = max(2, int(math.ceil(lengths_scaled[i] * 0.5)))

        # Generate heatmap
        if self.heatmap_encoding == "umich":
            if self.labels_encoding == "circle_ignore" and self.ignore_low_confidence_detections:
                draw_circle(heatmap[0], center_int, heatmap_radius, self.ignore_value)
            else:
                draw_umich_gaussian(heatmap[0], center_int, heatmap_radius)
        elif self.heatmap_encoding == "point":
            if self.labels_encoding == "circle_ignore" and self.ignore_low_confidence_detections:
                heatmap[0, center_int[1], center_int[0]] = self.ignore_value
            else:
                heatmap[0, center_int[1], center_int[0]] = 1.0
        else:
            raise KeyError(f"Unknown heatmap_encoding: {self.heatmap_encoding}")

        # Encode vessel and fishing maps
        if self.labels_encoding in {"circle", "circle_ignore"}:
            labels_radius = self.labels_radius if self.labels_radius else heatmap_radius
            draw_circle(is_vessel_map[0], center_int, labels_radius, self.vessel_mapping.get(is_vessel_sorted[i], self.ignore_value))
            draw_circle(is_fishing_map[0], center_int, labels_radius, self.fishing_mapping.get(is_fishing_sorted[i], self.ignore_value))
        elif self.labels_encoding == "point":
            # Handle point encoding if necessary
            raise NotImplementedError(f"Labels encoding 'point' not implemented")
        else:
            raise KeyError(f"Unknown labels_encoding: {self.labels_encoding}")

        # Encode lengths
        length_map[0, center_int[1], center_int[0]] = log_lengths[i]

        # Encode offset
        offset_map[0, center_int[1], center_int[0]] = center[0] - center_int[0]
        offset_map[1, center_int[1], center_int[0]] = center[1] - center_int[1]

    # Generate YOLO targets
    yolo_targets = self.generate_yolo_targets(centers_sorted, confidences_sorted, is_vessel_sorted, is_fishing_sorted, lengths_sorted)
    yolo_targets = yolo_targets.transpose(2, 0, 1)
    yolo_targets = torch.from_numpy(yolo_targets).float()  # Convert to torch.Tensor
    yolo_targets = yolo_targets.unsqueeze(0)   
                # Add batch dimension if necessary
    return MultilabelCircleNetEncodeResult(
        heatmap=heatmap,
        is_fishing=is_fishing_map,            
        is_vessel=is_vessel_map,
        lengths=length_map,
        offset=offset_map,
        yolo_targets=yolo_targets,  # Include yolo_targets in the result
    )

def generate_yolo_targets(
    self,
    centers: np.ndarray,
    confidences: np.ndarray,
    is_vessel: np.ndarray,
    is_fishing: np.ndarray,
    lengths: np.ndarray
) -> np.ndarray:
    
    num_objects = len(centers)
    yolo_targets = np.zeros((self.grid_height, self.grid_width, 6), dtype=np.float32)

    for i in range(num_objects):
        center = centers[i]
        confidence = self.confidence_mapping.get(confidences[i], self.ignore_value)  # Convert confidence to numerical value
        vessel_type = self.vessel_mapping.get(is_vessel[i], self.ignore_value)  # Ensure correct mapping
        fishing_status = self.fishing_mapping.get(is_fishing[i], self.ignore_value)  # Ensure correct mapping
        length = lengths[i]

        # Convert center coordinates to grid cell format (relative to the grid)
        grid_x = int(center[0] // self.output_stride)
        grid_y = int(center[1] // self.output_stride)

        # Ensure grid indices are within bounds
        grid_x = min(grid_x, self.grid_width - 1)
        grid_y = min(grid_y, self.grid_height - 1)

        # Assign to the correct grid cell
        yolo_targets[grid_y, grid_x, 0:4] = np.array([center[0], center[1], length, length])
        yolo_targets[grid_y, grid_x, 4] = confidence
        yolo_targets[grid_y, grid_x, 5] = vessel_type

    return yolo_targets



@torch.no_grad()
def decode(
    self,
    objectness_map: Tensor,
    is_vessel_map: Tensor,
    is_fishing_map: Tensor,
    length_map: Tensor,
    offset_map: Tensor,
    max_objects: Optional[int] = None,
    apply_activation: bool = False,
) -> MultilabelCircleNetDecodeResult:
   
    batch, _, height, width = objectness_map.size()

    objectness_map = cast_to_float_if_on_cpu(objectness_map)
    is_fishing_map = cast_to_float_if_on_cpu(is_fishing_map)
    is_vessel_map = cast_to_float_if_on_cpu(is_vessel_map)

    length_map = cast_to_float_if_on_cpu(length_map)

    if apply_activation:
        objectness_map = objectness_map.sigmoid()
        is_fishing_map = is_fishing_map.sigmoid()
        is_vessel_map = is_vessel_map.sigmoid()

    if max_objects is None:
        max_objects = self.max_objects

    # Handle non-masked images (with NaNs)
    objectness_map = torch.masked_fill(objectness_map, ~torch.isfinite(objectness_map), 0)
    is_fishing_map = torch.masked_fill(is_fishing_map, ~torch.isfinite(is_fishing_map), 0)
    is_vessel_map = torch.masked_fill(is_vessel_map, ~torch.isfinite(is_vessel_map), 0)

    obj_map = self.nms_method(objectness_map, kernel=self.nms_kernel_size)

    # Limit K to prevent having K more than W * H
    max_objects = min(max_objects, height * width)

    # Decode centers
    obj_scores, inds, _, ys, xs = centernet_topk(obj_map, top_k=max_objects)
    if offset_map is not None:
        offset_map = cast_to_float_if_on_cpu(offset_map)
        offset_map = centernet_tranpose_and_gather_feat(offset_map, inds)
        xs = xs.view(batch, max_objects, 1) + offset_map[:, :, 0:1]
        ys = ys.view(batch, max_objects, 1) + offset_map[:, :, 1:2]
    elif self.output_stride == 1:
        xs = xs.view(batch, max_objects, 1)
        ys = ys.view(batch, max_objects, 1)
    else:
        xs = xs.view(batch, max_objects, 1) + 0.5
        ys = ys.view(batch, max_objects, 1) + 0.5
    pred_centers = torch.cat([xs, ys], dim=2) * self.output_stride

    # Decode lengths
    pred_lengths = centernet_tranpose_and_gather_feat(length_map, inds)
    pred_lengths = pred_lengths[:, :, 0]  # Squeeze length

    # Decode is_vessel
    is_vessel = torch.sigmoid(centernet_tranpose_and_gather_feat(is_vessel_map, inds))  # Apply sigmoid for binary classification
    is_vessel = is_vessel.max(dim=2)[0]

    # Decode is_fishing
    is_fishing = torch.sigmoid(centernet_tranpose_and_gather_feat(is_fishing_map, inds))  # Apply sigmoid for binary classification
    is_fishing = is_fishing.max(dim=2)[0]

    return MultilabelCircleNetDecodeResult(
        centers=pred_centers,
        lengths=self.log2length(pred_lengths),  # Since lengths are in log-space, convert back to linear units
        is_fishing=is_fishing,
        is_vessel=is_vessel,
        scores=obj_scores,
    )

@classmethod
def from_config(cls, config: Dict[str, Any]) -> 'MultilabelCircleNetCoder':
    
    nms_method_name = config['nms_method']
    nms_method = NMS_METHODS.get(nms_method_name)
    
    if nms_method is None:
        raise ValueError(f"Unknown nms_method: {nms_method_name}")
    
    return cls(
        image_size=config['image_size'],
        output_stride=config['output_stride'],
        max_objects=config['max_objects'],
        heatmap_encoding=config['heatmap_encoding'],
        labels_encoding=config['labels_encoding'],
        ignore_value=config['ignore_value'],
        fixed_radius=config.get('fixed_radius'),
        labels_radius=config.get('labels_radius'),
        nms_kernel_size=config.get('nms_kernel_size', 3),
        nms_method=nms_method,
        ignore_low_confidence_detections=config.get('ignore_low_confidence_detections', False),
    )

@classmethod
def log2length(cls, x: Union[Tensor, np.ndarray]) -> Union[Tensor, np.ndarray]:
    
    if torch.is_tensor(x):
        return torch.exp(F.relu(x)) - 1
    else:
        return np.exp(np.clip(x, a_min=0, a_max=None)) - 1

@classmethod
def length2log(cls, x: np.ndarray) -> np.ndarray:
    
    return np.log(x + 1)
###needed review multilabel_dataset.py
import os
from functools import partial
from typing import List, Tuple, Optional, Callable, Dict, Any, Union
from xview3.centernet.models.heads.yolo_head import generate_yolo_targets
import albumentations as A
import cv2
import numpy as np
import pandas as pd
import torch
from torch import nn
from omegaconf import DictConfig
from pytorch_toolbelt.datasets import INPUT_IMAGE_KEY, INPUT_INDEX_KEY
from pytorch_toolbelt.utils.torch_utils import image_to_tensor
from sklearn.neighbors import BallTree
from sklearn.utils import compute_sample_weight
from torch.utils.data import Sampler, WeightedRandomSampler
from torch.utils.data.dataloader import default_collate
from tqdm import tqdm
from omegaconf import OmegaConf
from .bboxer import (
MultilabelCircleNetCoder,
MultilabelCircleNetEncodeResult,
)
from .constants import *
from ..constants import *
all = [
"MultilabelCircleNetDataModule",
"MultilabelCircleNetFixedCropDataset",
"MultilabelCircleNetRandomCropDataset",
"centernet_collate",
]
from ..dataset import (
RandomCropFromImageDataset,
RandomCropAroundEachObjectDataset,
XView3DataModule,
KeypointsDataset,
FixedCropFromImageDataset,
read_multichannel_image,
)
from ..dataset.random_crop_dataset import encode_labels
from ..metric import get_shoreline_shoreline_contours
def centernet_collate(batch, channels_last=False):
skip_keys = [
TARGET_CENTERS_KEY,
TARGET_FISHING_KEY,
TARGET_VESSEL_KEY,
TARGET_LABELS_KEY,
TARGET_LENGTHS_KEY,
INPUT_SCENE_ID_KEY,
INPUT_SCENE_CROP_KEY,
INPUT_CONFIDENCE_KEY,
INPUT_LOCATION_ID_KEY,
INPUT_IS_NEAR_SHORE_KEY,
]

excluded_items = {k: [] for k in skip_keys}
included_items = []

for b in batch:
    included = {k: v for k, v in b.items() if k not in skip_keys}
    included_items.append(included)
    for k in skip_keys:
        if k in b:
            excluded_items[k].append(b[k])

# Use default_collate for included items (numerical tensors)
batch_dict = default_collate(included_items)

# Process excluded items (categorical lists)
for k, v in excluded_items.items():
    if v:
        # Keep categorical fields as lists of lists or lists of tuples
        batch_dict[k] = v  # List of lists or list of tuples

# Handle channels_last if necessary
if channels_last and INPUT_IMAGE_KEY in batch_dict:
    batch_dict[INPUT_IMAGE_KEY] = batch_dict[INPUT_IMAGE_KEY].to(memory_format=torch.channels_last)

return batch_dict
class MultilabelCircleNetTargetsMixin:
box_coder: MultilabelCircleNetCoder

def __init__(
    self,
    stride=32,
    confidence_mapping=None,
    vessel_mapping=None,
    fishing_mapping=None,
    ignore_value=-1,
):
    self.stride = stride
    self.confidence_mapping = {
        'NONE': 0.0,
        'LOW': 0.33,
        'MEDIUM': 0.66,
        'HIGH': 1.0,
    }
    self.vessel_mapping = vessel_mapping or {0: 0, 1: 1}
    self.fishing_mapping = fishing_mapping or {0: 0, 1: 1}
    self.ignore_value = ignore_value


def generate_yolo_targets(self, centers, confidences, is_vessel, is_fishing, lengths):
  
    num_objects = len(centers)
    yolo_targets = np.zeros((self.grid_height, self.grid_width, 6), dtype=np.float32)

    for i in range(num_objects):
        center = centers[i]
        confidence_label = confidences[i]
        confidence = self.confidence_mapping[confidence_label]  # Map string label to numerical value
        vessel_type = self.vessel_mapping.get(is_vessel[i], self.ignore_value)
        fishing_status = self.fishing_mapping.get(is_fishing[i], self.ignore_value)
        length = lengths[i]

        # Compute grid indices
        grid_x = int(center[0] * self.grid_width)
        grid_y = int(center[1] * self.grid_height)

        # Ensure grid indices are within bounds
        grid_x = min(max(grid_x, 0), self.grid_width - 1)
        grid_y = min(max(grid_y, 0), self.grid_height - 1)

        # Assign the values
        yolo_targets[grid_y, grid_x, 0:4] = np.array([center[0], center[1], length, length])
        yolo_targets[grid_y, grid_x, 4] = confidence
        yolo_targets[grid_y, grid_x, 5] = vessel_type

    return yolo_targets

def build_centernet_sample(
    self,
    image,
    centers,
    confidences,
    lengths,
    is_vessel,
    is_fishing,
    index,
    image_id,
    is_near_shore,
    location,
    folder: str,
    crop_coords: Tuple[Tuple[int, int], Tuple[int, int]],
) -> Dict[str, Any]:
    # Mask NaN values
    image[~np.isfinite(image)] = 0
    box_coder = self.box_coder.create_box_coder_for_image_size(image.shape[:2])
    
    # Define grid dimensions based on image size and stride
    image_height, image_width = image.shape[:2]
    grid_height = image_height // self.box_coder.output_stride
    grid_width = image_width // self.box_coder.output_stride

    # Now use grid_height and grid_width
    yolo_targets = self.box_coder.generate_yolo_targets(
        centers,
        confidences,
        is_vessel,
        is_fishing,
        lengths,
    )

    # Convert numerical fields to NumPy arrays with explicit dtypes
    centers_np = np.array(centers, dtype=np.float32)
    lengths_np = np.array(lengths, dtype=np.float32)
    is_vessel_np = np.array(is_vessel, dtype=np.float32)
    is_fishing_np = np.array(is_fishing, dtype=np.float32)
    is_near_shore_np = np.array(is_near_shore, dtype=np.float32)
    #confidences_np = np.array(confidences, dtype=np.float32)

    # Validate consistency in lengths
    num_objects = len(confidences)
    if not all([
        len(centers_np) == num_objects,
        len(lengths_np) == num_objects,
        len(is_vessel_np) == num_objects,
        len(is_fishing_np) == num_objects,
    ]):
        raise ValueError("Mismatch in number of objects across targets")

    # Construct the result dictionary with tensors for numerical fields and lists for categorical fields
    result = {
        INPUT_INDEX_KEY: index,
        INPUT_IMAGE_KEY: image_to_tensor(image),  # Keeps image as tensor
        TARGET_CENTERS_KEY: torch.from_numpy(centers_np),  # Convert to tensor
        TARGET_VESSEL_KEY: torch.from_numpy(is_vessel_np),  # Convert to tensor
        TARGET_FISHING_KEY: torch.from_numpy(is_fishing_np),  # Convert to tensor
        TARGET_LENGTHS_KEY: torch.from_numpy(lengths_np),  # Convert to tensor
        INPUT_CONFIDENCE_KEY: confidences,  
        INPUT_IS_NEAR_SHORE_KEY: torch.from_numpy(is_near_shore_np),  # Convert to tensor
        INPUT_LOCATION_ID_KEY: location,
        INPUT_FOLDER_KEY: folder,
        INPUT_SCENE_ID_KEY: image_id,
        TARGET_SAMPLE_WEIGHT: {"train": 0.25, "validation": 1.0, "valid": 1.0, "test": 0.1}[folder],
    }
    if crop_coords is not None:
        result[INPUT_SCENE_CROP_KEY] = crop_coords  # Keep as tuple of tuples
    # Encode targets using the box coder
    centernet_targets: MultilabelCircleNetEncodeResult = box_coder.encode(
        centers=centers_np,
        confidences=confidences,  # Pass as list of categorical labels
        lengths=lengths_np,
        is_vessel=is_vessel_np,
        is_fishing=is_fishing_np,
    )

    # Convert encoded targets to PyTorch tensors
    result[CENTERNET_TARGET_OBJECTNESS_MAP] = torch.from_numpy(centernet_targets.heatmap).float()
    result[CENTERNET_TARGET_VESSEL_MAP] = torch.from_numpy(centernet_targets.is_vessel).float()
    result[CENTERNET_TARGET_FISHING_MAP] = torch.from_numpy(centernet_targets.is_fishing).float()
    result[CENTERNET_TARGET_SIZE] = torch.from_numpy(centernet_targets.lengths).float()
    result[CENTERNET_TARGET_OFFSET] = torch.from_numpy(centernet_targets.offset).float()
    result[CENTERNET_TARGET_YOLO] = torch.from_numpy(yolo_targets).float()
    

    return result
class MultilabelCircleNetFixedCropDataset(FixedCropFromImageDataset, MultilabelCircleNetTargetsMixin):
def init(
self,
scenes: List[str],
centers: List[np.ndarray],
locations: List[str],
confidences: List[np.ndarray],
is_vessel: List[np.ndarray],
is_fishing: List[np.ndarray],
is_near_shore: List[np.ndarray],
lengths: List[np.ndarray],
crop_coords: List[Tuple],
scene_sizes: List[Tuple[int, int]],
individual_transforms: Dict[str, A.Compose],
normalization: Dict[str, A.ImageOnlyTransform],
input_channels: List[str],
transform: A.Compose,
read_image_fn: Callable,
box_coder: MultilabelCircleNetCoder,
channels_last: bool = False,
):
super().init(
scenes=scenes,
centers=centers,
locations=locations,
confidences=confidences,
is_vessel=is_vessel,
is_fishing=is_fishing,
is_near_shore=is_near_shore,
lengths=lengths,
crop_coords=crop_coords,
scene_sizes=scene_sizes,
individual_transforms=individual_transforms,
normalization=normalization,
input_channels=input_channels,
transform=transform,
read_image_fn=read_image_fn,
channels_last=channels_last,
)
self.box_coder = box_coder

def build_sample(
    self,
    image,
    centers,
    confidences,
    is_vessel,
    is_fishing,
    lengths,
    index,
    folder: str,
    image_id,
    is_near_shore,
    location: str,
    crop_coords,
) -> Dict[str, Any]:
    box_coder = self.box_coder.create_box_coder_for_image_size(image.shape[:2])
    return self.build_centernet_sample(
        image=image,
        centers=centers,
        confidences=confidences,
        is_vessel=is_vessel,
        is_fishing=is_fishing,
        lengths=lengths,
        index=index,
        image_id=image_id,
        is_near_shore=is_near_shore,
        location=location,
        folder=folder,
        crop_coords=crop_coords,
    )

def get_collate_fn(self) -> Callable:
    return partial(centernet_collate, channels_last=self.channels_last)
class MultilabelCircleNetRandomCropDataset(RandomCropFromImageDataset, MultilabelCircleNetTargetsMixin):
def init(
self,
scenes: List[str],
shoreline_contours: List[np.ndarray],
locations: List[str],
centers: List[np.ndarray],
confidences: List[np.ndarray],
lengths: List[np.ndarray],
is_vessel: List[np.ndarray],
is_fishing: List[np.ndarray],
is_near_shore: List[np.ndarray],
scene_sizes: List[Tuple[int, int]],
individual_transforms: Dict[str, A.Compose],
normalization: Dict[str, A.ImageOnlyTransform],
input_channels: List[str],
transform: A.Compose,
read_image_fn: Callable,
crop_size: Tuple[int, int],
box_coder: MultilabelCircleNetCoder,
crop_around_ship_p=0.9,
balance_near_shore: Union[None, bool, float] = None,
balance_crowd=None,
balance_by_type=False,
channels_last=False,
):
super().init(
scenes=scenes,
scene_sizes=scene_sizes,
shoreline_contours=shoreline_contours,
locations=locations,
centers=centers,
confidences=confidences,
is_vessel=is_vessel,
is_fishing=is_fishing,
is_near_shore=is_near_shore,
lengths=lengths,
crop_size=crop_size,
transform=transform,
individual_transforms=individual_transforms,
normalization=normalization,
input_channels=input_channels,
read_image_fn=read_image_fn,
balance_near_shore=balance_near_shore,
balance_crowd=balance_crowd,
balance_by_type=balance_by_type,
channels_last=channels_last,
)
self.box_coder = box_coder

def build_sample(
    self,
    image,
    centers,
    confidences,
    is_vessel,
    is_fishing,
    lengths,
    index,
    folder: str,
    image_id,
    is_near_shore,
    location: str,
    crop_coords,
) -> Dict[str, Any]:
   
    box_coder = self.box_coder.create_box_coder_for_image_size(image.shape[:2])
    return self.build_centernet_sample(
        image=image,
        centers=centers,
        confidences=confidences,
        is_vessel=is_vessel,
        is_fishing=is_fishing,
        lengths=lengths,
        index=index,
        image_id=image_id,
        is_near_shore=is_near_shore,
        location=location,
        folder=folder,
        crop_coords=crop_coords,
    )

def get_collate_fn(self) -> Callable:
    return partial(centernet_collate, channels_last=self.channels_last)
class MultilabelCircleNetRandomCropAroundEachObjectDataset(RandomCropAroundEachObjectDataset, MultilabelCircleNetTargetsMixin):
def init(
self,
scenes: List[str],
locations: List[str],
crop_centers: np.ndarray,
centers: List[np.ndarray],
confidences: List[np.ndarray],
lengths: List[np.ndarray],
is_vessel: List[np.ndarray],
is_fishing: List[np.ndarray],
is_near_shore: List[np.ndarray],
scene_sizes: List[Tuple[int, int]],
individual_transforms: Dict[str, A.Compose],
normalization: Dict[str, A.ImageOnlyTransform],
input_channels: List[str],
transform: A.Compose,
read_image_fn: Callable,
crop_size: Tuple[int, int],
box_coder: MultilabelCircleNetCoder,
channels_last=False,
):
super().init(
scenes=scenes,
locations=locations,
crop_centers=crop_centers,
centers=centers,
confidences=confidences,
is_vessel=is_vessel,
is_fishing=is_fishing,
is_near_shore=is_near_shore,
lengths=lengths,
scene_sizes=scene_sizes,
individual_transforms=individual_transforms,
normalization=normalization,
input_channels=input_channels,
transform=transform,
read_image_fn=read_image_fn,
crop_size=crop_size,
channels_last=channels_last,
)
self.box_coder = box_coder

def build_sample(
    self,
    image,
    centers,
    confidences,
    is_vessel,
    is_fishing,
    lengths,
    index,
    folder: str,
    image_id,
    is_near_shore,
    location: str,
    crop_coords,
) -> Dict[str, Any]:
    box_coder = self.box_coder.create_box_coder_for_image_size(image.shape[:2])
    return self.build_centernet_sample(
        image=image,
        centers=centers,
        confidences=confidences,
        is_vessel=is_vessel,
        is_fishing=is_fishing,
        lengths=lengths,
        index=index,
        image_id=image_id,
        is_near_shore=is_near_shore,
        location=location,
        folder=folder,
        crop_coords=crop_coords,
    )

def get_collate_fn(self) -> Callable:
    return partial(centernet_collate, channels_last=self.channels_last)
class MultilabelCircleNetDataModule(XView3DataModule):
def init(self, config: DictConfig):
super().init(data_dir=config.dataset.data_dir)
self.config = config

@classmethod
def get_keypoint_params(cls):
    return A.KeypointParams(format="xys", label_fields=["confidences", "is_vessel", "is_fishing", "is_near_shore"])

@classmethod
def get_random_crop_each_object_dataset(
    cls,
    train_df: pd.DataFrame,
    train_image_size: Tuple[int, int],
    input_channels: List[str],
    box_coder_config: Dict,
    normalization: Dict[str, A.ImageOnlyTransform],
    num_samples: int = 1024,
    individual_augmentations: Optional[Dict] = None,
    augmentations: Optional[List[A.BasicTransform]] = None,
    balance_crowd=False,
    balance_per_scene=True,
    balance_near_shore=None,
    balance_vessel_type=False,
    balance_fishing_type=False,
    balance_by_type=False,
    balance_location=False,
    channels_last=False,
    read_image_fn=read_multichannel_image,
) -> Tuple[KeypointsDataset, Optional[Sampler]]:
    scene_sizes = dict(
        [
            (row["scene_id"], (row["rows"], row["cols"]))
            for i, row in pd.read_csv(os.path.join(os.path.dirname(__file__), "..", "scene_sizes.csv")).iterrows()
        ]
    )

    all_images = []
    all_coordinates = []
    all_confidences = []
    all_lengths = []
    all_size = []
    all_is_vessel = []
    all_is_fishing = []
    all_is_near_shore = []

    all_scenes = []
    all_locations = []

    center_near_shore = []
    center_neighbors = []
    center_coordinates = []
    center_is_vessel = []
    center_is_fishing = []
    center_is_near_shore = []

    for scene_id in tqdm(train_df.scene_id.unique(), desc="Preparing training dataset"):
        # Get detections only for current scene
        scene_df = train_df[train_df.scene_id == scene_id]
        num_instances_in_scene = len(scene_df)
        scene_size = scene_sizes[scene_id]
        scene_path = scene_df.scene_path.values[0]
        location_id = scene_df.location.values[0]

        targets = cls.get_multilabel_targets_from_df(scene_df)
        distance_from_shore_km = scene_df.distance_from_shore_km
        near_shore = (distance_from_shore_km < 2).values

        tree = BallTree(targets.centers, leaf_size=8, metric="chebyshev")
        neighbors = tree.query_radius(targets.centers, r=max(train_image_size) // 2, count_only=True)

        all_scenes.extend([scene_id] * num_instances_in_scene)
        all_locations.extend([location_id] * num_instances_in_scene)
        all_images.extend([scene_path] * num_instances_in_scene)

        center_coordinates.extend(targets.centers)
        center_is_vessel.extend(targets.is_vessel)
        center_is_fishing.extend(targets.is_fishing)
        center_is_near_shore.extend(targets.is_near_shore)
        center_near_shore.extend(near_shore)
        center_neighbors.extend(neighbors)

        all_coordinates.extend([targets.centers] * num_instances_in_scene)
        all_confidences.extend([targets.confidences] * num_instances_in_scene)
        all_is_vessel.extend([targets.is_vessel] * num_instances_in_scene)
        all_is_fishing.extend([targets.is_fishing] * num_instances_in_scene)
        all_lengths.extend([targets.lengths] * num_instances_in_scene)
        all_is_near_shore.extend([targets.is_near_shore] * num_instances_in_scene)

        all_size.extend([scene_size] * num_instances_in_scene)

    if augmentations is None:
        augmentations = []

    transform = A.Compose(
        augmentations
        + [
            A.PadIfNeeded(
                train_image_size[0],
                train_image_size[1],
                value=float("nan"),
                border_mode=cv2.BORDER_CONSTANT,
            ),
        ],
        keypoint_params=cls.get_keypoint_params(),
    )
    # Instantiate the box_coder using the configuration
    box_coder = MultilabelCircleNetCoder.from_config(box_coder_config)
    train_ds = MultilabelCircleNetRandomCropAroundEachObjectDataset(
        scenes=all_images,
        locations=all_locations,
        centers=all_coordinates,
        confidences=all_confidences,
        crop_centers=np.array(center_coordinates),
        lengths=all_lengths,
        is_vessel=all_is_vessel,
        is_fishing=all_is_fishing,
        is_near_shore=all_is_near_shore,
        scene_sizes=all_size,
        individual_transforms=individual_augmentations,
        input_channels=input_channels,
        transform=transform,
        normalization=normalization,
        read_image_fn=read_image_fn,
        box_coder=box_coder,
        crop_size=train_image_size,
        channels_last=channels_last,
    )

    weights = np.ones(len(all_coordinates), dtype=np.float64)
    num_balancing = 0

    if balance_location:
        location_balance_weights = compute_sample_weight("balanced", all_locations)
        weights += location_balance_weights
        num_balancing += 1

    if balance_crowd:
        inv_frequency = np.reciprocal(np.array(center_neighbors).astype(np.float32))
        crowd_balance_weights = inv_frequency
        if not np.all(np.isfinite(crowd_balance_weights)):
            raise RuntimeError()
        weights += crowd_balance_weights
        num_balancing += 1

    if balance_per_scene:
        scene_balance_weights = compute_sample_weight("balanced", all_scenes)
        weights += scene_balance_weights
        num_balancing += 1

    if balance_vessel_type:
        center_is_vessel = np.array(center_is_vessel)
        vessel_type_weights = compute_sample_weight(
            {
                0: len(center_is_vessel) / (2.0 * np.sum(center_is_vessel == 0)),
                1: len(center_is_vessel) / (2.0 * np.sum(center_is_vessel == 1)),
                IGNORE_LABEL: 0,
            },
            center_is_vessel,
        )
        if not np.all(np.isfinite(vessel_type_weights)):
            raise RuntimeError()

        weights += vessel_type_weights
        num_balancing += 1

    if balance_fishing_type:
        center_is_fishing = np.array(center_is_fishing)
        fishing_weights = compute_sample_weight(
            {
                0: len(center_is_fishing) / (2.0 * np.sum(center_is_fishing == 0)),
                1: len(center_is_fishing) / (2.0 * np.sum(center_is_fishing == 1)),
                IGNORE_LABEL: 0,
            },
            center_is_fishing,
        )
        if not np.all(np.isfinite(fishing_weights)):
            raise RuntimeError()
        weights += fishing_weights
        num_balancing += 1

    if balance_by_type:
        class_labels = encode_labels(center_is_vessel, center_is_fishing)
        type_balance_weights = compute_sample_weight("balanced", class_labels)
        weights += type_balance_weights
        num_balancing += 1

    if balance_near_shore and len(np.bincount(center_near_shore)) > 1:
        shore_balance_weights = compute_sample_weight("balanced", center_near_shore)
        weights += shore_balance_weights
        num_balancing += 1

    if num_balancing:
        weights /= float(num_balancing)

    train_sampler = WeightedRandomSampler(weights.astype(np.float32), num_samples=num_samples, replacement=True)
    return train_ds, train_sampler

@classmethod
def get_random_crop_training_dataset(
    cls,
    train_df: pd.DataFrame,
    train_image_size: Tuple[int, int],
    input_channels: List[str],
    box_coder: MultilabelCircleNetCoder,  
    num_samples: int = 1024,
    crop_around_ship_p: float = 0.9,
    individual_augmentations: Optional[Dict] = None,
    augmentations: Optional[Union[List[A.BasicTransform], A.Compose]] = None,
    #normalization: Dict[str, A.ImageOnlyTransform] = None,
    normalization: Union[Dict[str, A.ImageOnlyTransform], Callable] = None,
    balance_near_shore: Optional[bool] = None,
    balance_crowd: bool = False,
    balance_location: bool = False,
    balance_folder: bool = False,
    channels_last: bool = False,
    shore_root: str = None,
    read_image_fn=read_multichannel_image,
) -> Tuple[KeypointsDataset, Optional[Sampler]]:
    if augmentations is None:
        augmentations = []
    elif isinstance(augmentations, A.Compose):
        augmentations = augmentations.transforms

    transform = A.Compose(
        augmentations
        + [
            A.PadIfNeeded(
                train_image_size[0],
                train_image_size[1],
                value=float("nan"),
                border_mode=cv2.BORDER_CONSTANT,
            ),
        ],
        keypoint_params=cls.get_keypoint_params(),
    )

    scene_sizes = dict(
        [
            (row["scene_id"], (row["rows"], row["cols"]))
            for i, row in pd.read_csv(os.path.join(os.path.dirname(__file__), "..", "scene_sizes.csv")).iterrows()
        ]
    )

    all_images = []
    all_centers = []
    all_confidences = []
    all_lengths = []
    all_size = []
    all_is_vessel = []
    all_is_fishing = []
    all_is_near_shore = []
    all_near_shore = []
    all_locations = []
    all_shoreline_contours = []
    all_folders = []

    for scene_id in tqdm(train_df.scene_id.unique(), desc="Preparing training dataset"):
        # Get detections only for current scene
        scene_df = train_df[train_df.scene_id == scene_id]
        scene_size = scene_sizes[scene_id]
        scene_path = scene_df.scene_path.values[0]
        location_id = scene_df.location.values[0]
        folder_id = os.path.basename(os.path.dirname(scene_path))

        shoreline_contours = get_shoreline_shoreline_contours(shore_root, scene_id)

        targets = cls.get_multilabel_targets_from_df(scene_df)
        distance_from_shore_km = scene_df.distance_from_shore_km
        near_shore = (distance_from_shore_km < 2).values

        all_images.append(scene_path)
        all_shoreline_contours.append(shoreline_contours)

        all_centers.append(targets.centers)
        all_confidences.append(targets.confidences)
        all_is_vessel.append(targets.is_vessel)
        all_is_fishing.append(targets.is_fishing)
        all_is_near_shore.append(targets.is_near_shore)
        all_lengths.append(targets.lengths)
        all_size.append(scene_size)
        all_near_shore.append(near_shore)

        all_locations.append(location_id)
        all_folders.append(folder_id)
    
    train_ds = MultilabelCircleNetRandomCropDataset(
        scenes=all_images,
        shoreline_contours=all_shoreline_contours,
        locations=all_locations,
        centers=all_centers,
        confidences=all_confidences,
        lengths=all_lengths,
        is_vessel=all_is_vessel,
        is_fishing=all_is_fishing,
        is_near_shore=all_is_near_shore,
        scene_sizes=all_size,
        individual_transforms=individual_augmentations,
        input_channels=input_channels,
        transform=transform,
        normalization=normalization,
        read_image_fn=read_image_fn,
        box_coder=box_coder,
        crop_size=train_image_size,
        crop_around_ship_p=crop_around_ship_p,
        balance_near_shore=balance_near_shore,
        balance_crowd=balance_crowd,
        balance_by_type=False,
        channels_last=channels_last,
    )

    weights = len(train_df)
    num_balancing = 0

    if balance_location:
        location_balance_weights = compute_sample_weight("balanced", all_locations)
        weights *= location_balance_weights
        num_balancing += 1

    if balance_folder:
        folder_balance_weights = compute_sample_weight("balanced", all_folders)
        weights *= folder_balance_weights
        num_balancing += 1

    if num_balancing:
        weights **= 1.0 / num_balancing

    train_sampler = WeightedRandomSampler(np.ones(len(train_ds)), num_samples=num_samples, replacement=True)
    return train_ds, train_sampler

@classmethod
def get_validation_dataset(
    cls,
    valid_df: pd.DataFrame,
    valid_crop_size: Tuple[int, int],
    input_channels: List[str],
    box_coder: MultilabelCircleNetCoder,
    normalization: Dict[str, A.ImageOnlyTransform],
    channels_last=False,
    read_image_fn=read_multichannel_image,
) -> KeypointsDataset:
    """
    For validation dataset we apply sliding crops
    :param valid_df:
    :param kwargs:
    :return:
    """

    all_images = []
    all_centers = []
    all_confidences = []
    all_lengths = []
    all_crop_coords = []
    all_is_vessel = []
    all_is_fishing = []
    all_scene_sizes = []
    all_locations = []
    all_is_near_shore = []

    xview3_dir = os.environ.get('XVIEW3_DIR') #, '/Volumes/Crucial/data/'

    scene_sizes = dict(
        [
            (row["scene_id"], (row["rows"], row["cols"]))
            for i, row in pd.read_csv(os.path.join(os.path.dirname(__file__), "..", "scene_sizes.csv")).iterrows()
        ]
    )

    for scene_id in tqdm(valid_df.scene_id.unique(), desc="Preparing validation dataset"):
        # Get detections only for current scene
        scene_df = valid_df[valid_df.scene_id == scene_id]
        scene_path = os.path.join(xview3_dir, 'full', 'validation', scene_id)
        scene_size = scene_sizes[scene_id]
        location = scene_df.location.values[0]

        targets = cls.get_multilabel_targets_from_df(scene_df)
        for crop_coords in cls.iterate_crops(scene_size, valid_crop_size, valid_crop_size):
            all_images.append(scene_path)
            all_locations.append(location)
            all_centers.append(targets.centers)
            all_confidences.append(targets.confidences)
            all_is_vessel.append(targets.is_vessel)
            all_is_fishing.append(targets.is_fishing)
            all_is_near_shore.append(targets.is_near_shore)
            all_lengths.append(targets.lengths)
            all_crop_coords.append(crop_coords)
            all_scene_sizes.append(scene_size)

    individual_transforms = {}
    transform = A.Compose(
        [
            A.PadIfNeeded(
                valid_crop_size[0],
                valid_crop_size[1],
                value=float("nan"),
                border_mode=cv2.BORDER_CONSTANT,
            ),
        ],
        keypoint_params=cls.get_keypoint_params(),
    )
    # Instantiate the box_coder using the configuration
    valid_ds = MultilabelCircleNetFixedCropDataset(
        scenes=all_images,
        locations=all_locations,
        centers=all_centers,
        confidences=all_confidences,
        lengths=all_lengths,
        crop_coords=all_crop_coords,
        is_vessel=all_is_vessel,
        is_fishing=all_is_fishing,
        is_near_shore=all_is_near_shore,
        scene_sizes=all_scene_sizes,
        individual_transforms=individual_transforms,
        normalization=normalization,
        input_channels=input_channels,
        transform=transform,
        read_image_fn=read_image_fn,
        box_coder=box_coder,
        channels_last=channels_last,
    )
    return valid_ds

I'm glad you've applied the suggestions and are seeking a review of your updated multilabel_circle_coder.py and multilabel_dataset.py. I'll carefully review your code to ensure compatibility after integrating YOLO with the multi-label dataset.
1. Review of multilabel_circle_coder.py

a. General Structure
You've added yolo_targets to the MultilabelCircleNetEncodeResult namedtuple and are generating YOLO targets within the encode method.
The generate_yolo_targets method is implemented within the MultilabelCircleNetCoder class.
b. encode Method
Let's look at the encode method, specifically where you generate YOLO targets:

# Generate YOLO targets
yolo_targets = self.generate_yolo_targets(centers_sorted, confidences_sorted, is_vessel_sorted, is_fishing_sorted, lengths_sorted)
yolo_targets = yolo_targets.transpose(2, 0, 1)
yolo_targets = torch.from_numpy(yolo_targets).float()  # Convert to torch.Tensor
yolo_targets = yolo_targets.unsqueeze(0)  # Add batch dimension if necessary
c. generate_yolo_targets Method
Your generate_yolo_targets method computes YOLO-style targets. Here's the method:

def generate_yolo_targets(
    self,
    centers: np.ndarray,
    confidences: np.ndarray,
    is_vessel: np.ndarray,
    is_fishing: np.ndarray,
    lengths: np.ndarray
) -> np.ndarray:
    # [Your implementation]
Issue Identified: Incorrect Computation of YOLO Targets
In YOLO, targets are defined somewhat differently from what you're currently assigning. Specifically, YOLO expects the following format for each grid cell that contains an object:
tx, ty: The position of the center of the object relative to the cell (values between 0 and 1).
tw, th: The width and height of the object, normalized by the image dimensions or anchor box dimensions.
confidence: Indicates if an object is present (usually 1) and sometimes includes IoU.
class probabilities: One-hot encoded or class indices.
In your code, you're assigning absolute coordinates and lengths directly to the grid cell, which is not aligned with YOLO's expected format.
Corrected Implementation:
Adjust the generate_yolo_targets method as follows:

def generate_yolo_targets(
    self,
    centers: np.ndarray,
    confidences: np.ndarray,
    is_vessel: np.ndarray,
    is_fishing: np.ndarray,
    lengths: np.ndarray
) -> np.ndarray:
    num_objects = len(centers)
    yolo_targets = np.zeros((self.grid_height, self.grid_width, 6), dtype=np.float32)

    for i in range(num_objects):
        center = centers[i]
        confidence_label = confidences[i]
        confidence = self.confidence_mapping.get(confidence_label, self.ignore_value)
        vessel_type = self.vessel_mapping.get(is_vessel[i], self.ignore_value)
        fishing_status = self.fishing_mapping.get(is_fishing[i], self.ignore_value)
        length = lengths[i]

        # Convert center coordinates to grid cell indices
        grid_x = int(center[0] // self.output_stride)
        grid_y = int(center[1] // self.output_stride)

        # Compute the offset of the center within the grid cell (values between 0 and 1)
        x_cell = (center[0] / self.output_stride) - grid_x
        y_cell = (center[1] / self.output_stride) - grid_y

        # Normalize width and height (lengths) by image dimensions (or anchor dimensions if using anchors)
        w_normalized = length / self.image_width
        h_normalized = length / self.image_height

        # Ensure grid indices are within bounds
        grid_x = min(grid_x, self.grid_width - 1)
        grid_y = min(grid_y, self.grid_height - 1)

        # Assign to the correct grid cell
        yolo_targets[grid_y, grid_x, 0] = x_cell
        yolo_targets[grid_y, grid_x, 1] = y_cell
        yolo_targets[grid_y, grid_x, 2] = w_normalized
        yolo_targets[grid_y, grid_x, 3] = h_normalized
        yolo_targets[grid_y, grid_x, 4] = confidence
        yolo_targets[grid_y, grid_x, 5] = vessel_type  # Or class index if you have multiple classes
Explanation:
x_cell and y_cell: Offsets inside the grid cell, values between 0 and 1.
w_normalized and h_normalized: Width and height normalized by the image dimensions.
Grid Indices (grid_x, grid_y): Calculated by integer division of the center coordinates by the output_stride.
Assignments to yolo_targets: Now follow the YOLO convention.
d. Ensure Consistency in Class Indices
Make sure that vessel_type (or class index) is an integer starting from 0 up to num_classes - 1.
e. Update the Transpose and Squeeze
After generating the yolo_targets:

# Transpose to [channels, grid_height, grid_width]
yolo_targets = yolo_targets.transpose(2, 0, 1)  # Shape: [6, grid_height, grid_width]

# Convert to torch.Tensor and add batch dimension
yolo_targets = torch.from_numpy(yolo_targets).float()  # [6, H, W]
yolo_targets = yolo_targets.unsqueeze(0)  # [1, 6, H, W]
2. Review of multilabel_dataset.py

a. Redundant generate_yolo_targets Method
In your MultilabelCircleNetTargetsMixin class, you have a generate_yolo_targets method. However, since yolo_targets are generated within the MultilabelCircleNetCoder.encode method, you don't need to generate them again in the dataset class.
Recommendation:
Remove the generate_yolo_targets method from MultilabelCircleNetTargetsMixin to avoid redundancy and potential inconsistencies.
b. Simplify build_centernet_sample Method
In your build_centernet_sample method, you currently:
Call box_coder.encode(...) to get centernet_targets.
Separately generate yolo_targets using self.box_coder.generate_yolo_targets(...).
This leads to duplicate calculations and may result in inconsistencies if the two methods diverge.
Recommendation:
Use yolo_targets directly from centernet_targets.yolo_targets.
Adjusted build_centernet_sample Method:

def build_centernet_sample(
    self,
    image,
    centers,
    confidences,
    lengths,
    is_vessel,
    is_fishing,
    index,
    image_id,
    is_near_shore,
    location,
    folder: str,
    crop_coords: Tuple[Tuple[int, int], Tuple[int, int]],
) -> Dict[str, Any]:
    # Mask NaN values
    image[~np.isfinite(image)] = 0
    box_coder = self.box_coder.create_box_coder_for_image_size(image.shape[:2])
    
    # Convert numerical fields to NumPy arrays with explicit dtypes
    centers_np = np.array(centers, dtype=np.float32)
    lengths_np = np.array(lengths, dtype=np.float32)
    is_vessel_np = np.array(is_vessel, dtype=np.float32)
    is_fishing_np = np.array(is_fishing, dtype=np.float32)
    is_near_shore_np = np.array(is_near_shore, dtype=np.float32)

    # Validate consistency in lengths
    num_objects = len(centers_np)
    if not all([
        len(confidences) == num_objects,
        len(lengths_np) == num_objects,
        len(is_vessel_np) == num_objects,
        len(is_fishing_np) == num_objects,
    ]):
        raise ValueError("Mismatch in number of objects across targets")

    result = {
        INPUT_INDEX_KEY: index,
        INPUT_IMAGE_KEY: image_to_tensor(image),
        TARGET_CENTERS_KEY: torch.from_numpy(centers_np),
        TARGET_VESSEL_KEY: torch.from_numpy(is_vessel_np),
        TARGET_FISHING_KEY: torch.from_numpy(is_fishing_np),
        TARGET_LENGTHS_KEY: torch.from_numpy(lengths_np),
        INPUT_CONFIDENCE_KEY: confidences,  # Keep as list of strings
        INPUT_IS_NEAR_SHORE_KEY: torch.from_numpy(is_near_shore_np),
        INPUT_LOCATION_ID_KEY: location,
        INPUT_FOLDER_KEY: folder,
        INPUT_SCENE_ID_KEY: image_id,
        TARGET_SAMPLE_WEIGHT: {"train": 0.25, "validation": 1.0, "valid": 1.0, "test": 0.1}[folder],
    }
    if crop_coords is not None:
        result[INPUT_SCENE_CROP_KEY] = crop_coords

    # Encode targets using the box coder
    centernet_targets: MultilabelCircleNetEncodeResult = box_coder.encode(
        centers=centers_np,
        confidences=confidences,  # List of strings
        lengths=lengths_np,
        is_vessel=is_vessel,  # Use the original list or map to strings as needed
        is_fishing=is_fishing,
    )

    # Extract and assign the targets
    result[CENTERNET_TARGET_OBJECTNESS_MAP] = torch.from_numpy(centernet_targets.heatmap).float()
    result[CENTERNET_TARGET_VESSEL_MAP] = torch.from_numpy(centernet_targets.is_vessel).float()
    result[CENTERNET_TARGET_FISHING_MAP] = torch.from_numpy(centernet_targets.is_fishing).float()
    result[CENTERNET_TARGET_SIZE] = torch.from_numpy(centernet_targets.lengths).float()
    result[CENTERNET_TARGET_OFFSET] = torch.from_numpy(centernet_targets.offset).float()

    # YOLO targets
    result[CENTERNET_TARGET_YOLO] = centernet_targets.yolo_targets  # Already a torch.Tensor
    
    return result
Note:
Ensure that is_vessel and is_fishing align with the expected types in box_coder.encode. If they should be lists of strings, pass them as such.
Your box_coder.encode method expects is_vessel and is_fishing as arrays (or lists) of strings to map them correctly.
c. Update centernet_collate Function
In your centernet_collate function, make sure to handle CENTERNET_TARGET_YOLO appropriately.
Add CENTERNET_TARGET_YOLO to skip_keys if necessary:

def centernet_collate(batch, channels_last=False):
    skip_keys = [
        TARGET_CENTERS_KEY,
        TARGET_FISHING_KEY,
        TARGET_VESSEL_KEY,
        TARGET_LABELS_KEY,
        TARGET_LENGTHS_KEY,
        INPUT_SCENE_ID_KEY,
        INPUT_SCENE_CROP_KEY,
        INPUT_CONFIDENCE_KEY,
        INPUT_LOCATION_ID_KEY,
        INPUT_IS_NEAR_SHORE_KEY,
        CENTERNET_TARGET_YOLO,  # Add this if you need special handling
    ]
    # Rest of the function...
However, since CENTERNET_TARGET_YOLO is a tensor and should be collated by default_collate, you might not need to add it to skip_keys.
d. Ensure Data Consistency
Verify that all tensors have consistent shapes and data types.
Make sure that any mappings for is_vessel and is_fishing are consistent between the dataset and the coder.
3. Additional Recommendations

a. Remove Unnecessary Initializations in MultilabelCircleNetTargetsMixin
Since you no longer need the generate_yolo_targets method in MultilabelCircleNetTargetsMixin, you may simplify its __init__ method:

class MultilabelCircleNetTargetsMixin:
    box_coder: MultilabelCircleNetCoder

    def __init__(self):
        pass  # If no initialization is needed
Or remove the __init__ method altogether if it's not required.
b. Consistency in is_vessel and is_fishing
Ensure that is_vessel and is_fishing are consistently represented throughout your codebase. If they are categorical variables requiring mapping, make sure:
In the dataset, they are passed as lists of strings representing categories.
In the coder, you have appropriate mappings to convert these strings into numerical values.
c. Model Adjustments
Ensure that your model's output layers are compatible with the new yolo_targets.
Modify your model's architecture to include the YOLO head if necessary.
Adjust your loss functions to handle both CenterNet and YOLO outputs accordingly.
4. Testing and Validation

Finally, to ensure compatibility and correctness:
a. Unit Tests
Write unit tests for generate_yolo_targets to confirm that the targets are generated as expected.
Check that positions, sizes, and class labels are correctly assigned and normalized.
b. Data Loader Testing
Run a small batch through your data loader and print out shapes and sample values to verify that yolo_targets are correctly included.
c. End-to-End Testing
Perform a forward pass through your model to ensure that the outputs match the expected shapes.
Compute the loss using your updated loss function to verify that there are no runtime errors.
d. Visual Inspection
Visualize some of the generated targets (both CenterNet and YOLO) overlaid on the images to manually check correctness.
Summary

Adjust the generate_yolo_targets method in multilabel_circle_coder.py to correctly compute YOLO targets following the YOLO format.
Remove redundant generate_yolo_targets method from MultilabelCircleNetTargetsMixin in multilabel_dataset.py.
Simplify the build_centernet_sample method to use the yolo_targets from centernet_targets.
Ensure data consistency for is_vessel, is_fishing, and other mappings.
Update the centernet_collate function if necessary, but ensure tensors are correctly batched.
Test thoroughly at each step to validate your changes.
By making these adjustments, your code should be compatible and function correctly after integrating YOLO with your multi-label dataset. If you encounter any specific issues or have further questions, feel free to ask!