to be able to give the correct solution you need to see these files
1-config
2-centernet_unet.py
3-yolo_head.py
4-yolo_decoder.py
5-decoupled_head.py
6-unet_block.py
7-custom_encoder.py
8-necks.py
9-transformer_classifier.py
and before all you should understand that
This project is a complex integrated multi-model system designed for vessel detection and fishing activity classification using SAR (Synthetic Aperture Radar) TIFF images as input. The system leverages several advanced deep learning architectures and techniques to achieve high accuracy and efficiency in maritime surveillance and monitoring.
input csv labels train, validation
train
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
validation
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
44.28796374,13.04173288,,manual,9923,23071,True,,9999.99,264ed833a13b7f2av,LOW,9918.0,23068.0,9928.0,23074.0,264ed833a13b7f2av_044.28796374000000213300_013.04173287999999963915
44.61788761,12.35710138,,manual,6105,17727,True,,6.25220144977796,264ed833a13b7f2av,LOW,6099.0,17725.0,6111.0,17729.0,264ed833a13b7f2av_044.61788760999999681189_012.35710137999999957970
44.28847473,13.7108852,,manual,10023,28410,True,,9999.99,264ed833a13b7f2av,MEDIUM,10018.0,28408.0,10028.0,28412.0,264ed833a13b7f2av_044.28847472999999723697_013.71088519999999988386
Key Components
CenterNet UNet
Combines CenterNet for object detection and U-Net for image segmentation.
YOLO (You Only Look Once)
A real-time object detection system known for its speed and accuracy.
Neural Architecture Search (NAS)
Optimizes model architecture to enhance performance.
Feature Pyramid Network (FPN)
Enables multi-scale feature extraction, improving detection across various object sizes.
Feature Aggregation Network (FAN)
Combines features from different network levels for enriched representation.
Swin Transformer
A hierarchical vision transformer that captures global context effectively.
Optuna
A hyperparameter optimization framework to fine-tune model settings.
Output Format
The system processes SAR images and outputs predictions in a CSV format with the following columns:
vessel_length_m: Length of the detected vessel in meters.
detect_scene_row: Row coordinate in the detection scene.
detect_scene_column: Column coordinate in the detection scene.
scene_id: Identifier for the scene.
is_vessel: Boolean indicating if a vessel is detected.
is_fishing: Boolean indicating if fishing activity is detected.
Sample Output:
vessel_length_m,detect_scene_row,detect_scene_column,scene_id,is_vessel,is_fishing
33.99261,11448,13456,1,True,True
28.166428,15632,12622,1,True,True
32.700775,19454,4840,1,True,True
37.80931,12862,2608,1,True,True
48.392456,4902,5522,1,True,False
The system architecture encompasses model configuration, data augmentation, training processes, and inference pipelines
so please understand all this and remember and respond with done so I can show you the files to deeply review and remember and respond with only 'done' so we can process with find solution
###config
defaults:
self
dataset: full_4fold
augs: flips_light
sampler: balanced
normalization: default
loss: centernet/rfl_bce_mse
model: centernet/b5_unet_s2
optimizer: adamw
scheduler: cos
train: default
runner: default
ema: disabled
torch: default
slug: b5_unet_s2
activation: silu
num_channels: ${dataset.num_channels}
num_extra_blocks: 0
model:
target: xview3.centernet.models.centernet_unet.EnhancedCenterNetUNet
slug: b5_unet_s2
weights: IMAGENET1K_V1
pretrained: True
output_stride: 16
image_size: [2048, 2048]
window_size:
4
4
num_layers: [2, 5]
num_filters: [64, 128, 256]
kernel_size: [3, 5]
activation: silu
dropout_rate: [0.0, 0.5]
dropout_rate_max: 0.5
optimizer: adamw
learning_rate: 0.001
transformer:
enabled: true
embed_dim: 96
depths: [2, 2, 6, 2]
num_heads: [3, 6, 12, 24]
window_size:
4
4
mlp_ratio: 4.0
qkv_bias: true
qk_scale: null
drop_rate: 0.0
attn_drop_rate: 0.0
drop_path_rate: 0.1
norm_layer: torch.nn.LayerNorm
use_dropout: false
use_custom_pos_embed: false
decoder:
target: xview3.centernet.models.custom_encoder.DebugUnetDecoder
encoder_channels: [24, 48, 64, 160, 272] #[32, 24, 40, 112, 320]
decoder_channels: [512, 384, 256, 128, 64]
unet_block:
target: xview3.centernet.models.unet_blocks.UnetBlock
partial: true
upsample_block: xview3.centernet.models.custom_encoder.UpsampleBlock
#scale_factor: 2
#mode: bilinear
center: False # Changed from 'bilinear' to 'nearest'
activation: silu
encoder:
target: pytorch_toolbelt.modules.encoders.TimmEfficientNetV2
model_name: efficientnetv2_rw_s
pretrained: True #features_only: True
activation: silu
drop_rate: 0.1
drop_path_rate: 0.1
layers:
1 # 4
2 # 8
3 # 16
4 # 32
nas:
enabled: true
search_space:
num_layers: [2, 5]
num_filters: [64, 128, 256]
kernel_size: [3, 5]
activation: [relu, leaky_relu, silu]
dropout_rate: [0.0, 0.5]
dropout_rate_max: 0.5
num_trials: 50
metric: "val_acc"
direction: "maximize"
num_channels: ${dataset.num_channels}
num_classes: 2
anchors:
[10, 13]
[16, 30]
[33, 23]
[30, 61]
[62, 45]
[59, 119]
input_upsample_factor: 1
neck_out_channels: 256
box_coder:
target: xview3.centernet.bboxer.multilabel_circle_coder.MultilabelCircleNetCoder
image_size:
2048
2048
max_objects: 512
heatmap_encoding: umich
labels_encoding: circle
ignore_value:
target: xview3.ignore_value
fixed_radius: 3
labels_radius: 1
ignore_low_confidence_detections: ${dataset.ignore_low_confidence_detections}
head:
decoupled_head:
target: xview3.centernet.models.heads.DecoupledHeadGroupNormLateShuffle
channels: 256
num_classes: 2
classifier_dim: 128
objectness_dim: 128
size_dim: 128
offset_dim: 16
dropout_rate: 0.1
activation: silu
num_blocks: 3
anchors: [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45], [59, 119]]
yolo_head:
target: xview3.centernet.models.heads.yolo_head.YOLOHead
num_classes: 2
input_channels: 256
anchors: [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45], [59, 119]]
input_image_size: [2048, 2048]
stride: 16
num_anchors_per_scale: 3
fan:
target: xview3.centernet.models.necks.FeatureAggregationNetwork
input_channels: [272] #[32, 48, 64, 128] #in_channels: 272 #out_channels: 256 #num_layers: 2
out_channels: 256
num_layers: 2
upsample_mode: bilinear
use_attention: True
fusion_method: weighted_sum
expected_fan_input_channels: 256
Per-Channel Normalization
normalization:
channels:
vv:
target: xview3.SigmoidNormalization
midpoint: -20.0
temperature: 0.18
vh:
target: xview3.SigmoidNormalization
midpoint: -20.0
temperature: 0.18
bathymetry:
target: xview3.CubicRootNormalization
wind_direction:
target: albumentations.NoOp
wind_speed:
target: albumentations.NoOp
mask:
target: albumentations.NoOp
mean: [0.485, 0.456, 0.406]
std: [0.229, 0.224, 0.225]
Dataset Configuration
dataset:
normalization:
channels:
vv:
target: xview3.SigmoidNormalization
midpoint: -20.0
temperature: 0.18
vh:
target: xview3.SigmoidNormalization
midpoint: -20.0
temperature: 0.18
bathymetry:
target: xview3.CubicRootNormalization
wind_direction:
target: albumentations.NoOp
wind_speed:
target: albumentations.NoOp
mask:
target: albumentations.NoOp
mean: [0.485, 0.456, 0.406]
std: [0.229, 0.224, 0.225]
slug: full_4fold
splitter: full
fold: 0
num_folds: 4
Added the following parameters to address references in box_coder
ignore_low_confidence_detections: true
ignore_low_confidence_labels: false
filter_low_confidence_objects: False
num_channels: 2
input_channels:
vh
vv
DDP Stuff
local_rank: ${oc.env:LOCAL_RANK, 0}
world_size: ${oc.env:WORLD_SIZE, 0}
experiment: null
seed: 42
deterministic: False
detect_anomaly: False
find_unused: False
transfer: null
checkpoint: null
inference:
channels_last: False
use_traced_model: True
tile_size: 2048
tile_step: 1536
accumulate_on_gpu: False
batch_size: 1
fp16: True
Torch Configuration
torch:
amp:
enabled: false
cudnn:
allow_tf32: True
cuda_allow_tf32: True
detect_anomaly: False
deterministic: False
benchmark: True
Inference and Training Enhancements
activation_after: "after_model"
tta: "d4"
with_offset: True
yolo_decoders: []
Validation Configuration
valid:
image_size: [2048, 2048]
batch_size: 2
num_workers: 2
Training Configuration
train:
slug: default_training
epochs: 4
output_stride: 16
early_stopping:
enabled: true
patience: 2
mode: "max" # Options: 'min' for loss-based, 'max' for accuracy-based
metric: "loss" # The metric to monitor for early stopping
image_size: [2048, 2048]
batch_size: 2
num_workers: 2
num_samples: 8192
persistent_workers: False
model2:
target: xview3.centernet.models.centernet_unet.EnhancedCenterNetUNet
slug: b5_unet_s2
weights: IMAGENET1K_V1
pretrained: True
output_stride: 16
image_size: [2048, 2048]
window_size:
4
4
num_layers: ${model2.num_layers} # Ensure these are set up for overrides
num_filters: ${model2.num_filters}
kernel_size: ${model2.kernel_size}
activation: ${model2.activation}
dropout_rate: ${model2.dropout_rate}
dropout_rate_max: 0.5
optimizer: ${model2.optimizer}
learning_rate: ${model2.learning_rate}
transformer:
enabled: true
embed_dim: 96
depths: [2, 2, 6, 2]
num_heads: [3, 6, 12, 24]
window_size:
4
4
mlp_ratio: 4.0
qkv_bias: true
qk_scale: null
drop_rate: 0.0
attn_drop_rate: 0.0
drop_path_rate: 0.1
norm_layer: torch.nn.LayerNorm
use_dropout: false
use_custom_pos_embed: false
decoder:
target: xview3.centernet.models.custom_encoder.DebugUNetDecoder
encoder_channels: [24, 48, 64, 160, 272]
channels: [512, 384, 256, 128, 64]
unet_block:
target: xview3.centernet.models.unet_blocks.UnetBlock
partial: true
upsample_block:
target: torch.nn.Upsample
scale_factor: 2
mode: bilinear
encoder:
target: pytorch_toolbelt.modules.encoders.TimmEfficientNetV2
model_name: efficientnetv2_rw_s
pretrained: True
activation: silu
drop_rate: 0.1
drop_path_rate: 0.1
layers:
1 # 4
2 # 8
3 # 16
4 # 32
nas:
enabled: true
search_space:
num_layers: [2, 5]
num_filters: [64, 128, 256]
kernel_size: [3, 5]
activation: [relu, leaky_relu, silu]
dropout_rate: [0.0, 0.5]
dropout_rate_max: 0.5
num_trials: 50
metric: "val_acc"
direction: "maximize"
num_channels: ${dataset.num_channels}
num_classes: 2
anchors: [[10, 13], [16, 30], [33, 23]]
input_upsample_factor: 1
neck_out_channels: 256
box_coder:
target: xview3.centernet.bboxer.multilabel_circle_coder.MultilabelCircleNetCoder
image_size:
2048
2048
max_objects: 512
heatmap_encoding: umich
labels_encoding: circle
ignore_value:
target: xview3.ignore_value
fixed_radius: 3
labels_radius: 1
ignore_low_confidence_detections: ${dataset.ignore_low_confidence_detections}
head:
decoupled_head:
target: xview3.centernet.models.heads.DecoupledHeadGroupNormLateShuffle
channels: 256
classifier_dim: 128
objectness_dim: 128
size_dim: 128
offset_dim: 16
dropout_rate: 0.1
activation: silu
num_blocks: 3
yolo_head:
target: xview3.centernet.models.heads.yolo_head.YOLOHead
num_classes: 2
anchors: [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45], [59, 119]]
input_image_size: [2048, 2048]
stride: 16
num_anchors_per_scale: 3
fan:
target: xview3.centernet.models.necks.FeatureAggregationNetwork
input_channels: [272] #[32, 48, 64, 128] #in_channels: 272 #out_channels: 256 #num_layers: 2
out_channels: 256
num_layers: 2
upsample_mode: bilinear
use_attention: True
fusion_method: weighted_sum

done

perfect now I'm integrating many libraries with CenterNetUnet as FAN, NAS , FBN, Optuna, Swin transformer and many more also the most important is YOLO which I've build for it head as yolo_head.py and yolo_decoder.py so it can also detect the vessels and give outputs faster and more accurate than CenterNetUnet
here's the yolo head
##yolo_head.py
import torch.nn as nn
from torch import Tensor
import torch
from typing import List, Tuple, Union, Dict
from ..necks import EnhancedFPNFuse as FPNFuse
import numpy as np
from ...constants import CENTERNET_OUTPUT_YOLO
import logging
logger = logging.getLogger(name)
class AnchorFreeHead(nn.Module):
def init(self, in_channels, num_classes):
super().init()
self.cls_conv = nn.Conv2d(in_channels, num_classes, 1)
self.reg_conv = nn.Conv2d(in_channels, 4, 1)
self.obj_conv = nn.Conv2d(in_channels, 1, 1)

def forward(self, x):
    cls_pred = self.cls_conv(x)
    reg_pred = self.reg_conv(x)
    obj_pred = self.obj_conv(x)
    return torch.cat([reg_pred, obj_pred, cls_pred], dim=1)
class YOLOHead(nn.Module):
def init(
self,
num_classes: int,
anchors: List[Tuple[int, int]],
input_channels: int,
input_image_size=(512, 512),
stride: int = 16
):
super(YOLOHead, self).init()
self.num_classes = num_classes
self.anchors = anchors
self.num_anchors = len(anchors)
self.stride = stride
self.input_image_size = input_image_size
self.out_channels = self.num_anchors * (5 + self.num_classes)

    # Define the convolutional layer using input_channels
    self.yolo_conv = nn.Sequential(
        nn.Conv2d(input_channels, 128, kernel_size=3, padding=1, bias=False),
        nn.BatchNorm2d(128),
        nn.ReLU(inplace=True),
        nn.Conv2d(128, self.out_channels, kernel_size=1)
    )

def forward(self, x: Tensor) -> Dict[str, Tensor]:
    yolo_output = self.yolo_conv(x)
    yolo_output = nn.functional.interpolate(yolo_output, scale_factor=2, mode='bilinear', align_corners=False)
    logger.debug(f"YOLOHead output shape after upsampling: {yolo_output.shape}")
    return {
        CENTERNET_OUTPUT_YOLO: yolo_output
    }
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
###here's yolo_decoder.py
import torch
import torch.nn as nn
from torch import Tensor
import numpy as np
from typing import List, Tuple, Dict
from dataclasses import dataclass
from xview3.utils import TorchDevice
import logging
logger = logging.getLogger(name)
device = TorchDevice.choose_torch_device()
@dataclass
class YOLODecodeResult:
boxes: Tensor
scores: Tensor
classes: Tensor
class YOLODecoder(nn.Module):
def init(self,
anchors: List[Tuple[int, int]],
num_classes: int,
input_image_size: Tuple[int, int],
stride: int,
conf_threshold: float = 0.5,
iou_threshold: float = 0.4):
super(YOLODecoder, self).init()
self.register_buffer('anchors', torch.tensor(anchors, dtype=torch.float32))
self.num_classes = num_classes
self.input_image_size = input_image_size
self.stride = stride
self.conf_threshold = conf_threshold
self.iou_threshold = iou_threshold
self.num_anchors = len(self.anchors)
self.out_channels = self.num_anchors * (5 + self.num_classes)
self.grid_size = (input_image_size[0] // stride, input_image_size[1] // stride)

def forward(self, x: Tensor) -> Tensor:
    return self.conv(x)

def decode(self, yolo_output: Tensor) -> YOLODecodeResult:
    """
    Decodes YOLO head outputs into bounding boxes, scores, and class predictions.

    Args:
        yolo_output (Tensor): Output from YOLO head of shape [B, num_anchors*(5+num_classes), H, W].

    Returns:
        YOLODecodeResult: Decoded bounding boxes, scores, and class indices.
    """
    batch_size, _, grid_h, grid_w = yolo_output.shape
    num_anchors = self.num_anchors
    device = yolo_output.device

    # Verify total elements
    expected_num_elements = batch_size * num_anchors * (5 + self.num_classes) * grid_h * grid_w
    actual_num_elements = yolo_output.numel()
    if actual_num_elements != expected_num_elements:
        logger.warning(
            f"Expected {expected_num_elements} elements, but got {actual_num_elements}."
            f" Proceeding with decoding, but check for potential issues."
        )

    # Reshape and continue as before
    yolo_output = yolo_output.view(batch_size, num_anchors, 5 + self.num_classes, grid_h, grid_w)
    yolo_output = yolo_output.permute(0, 1, 3, 4, 2).contiguous()  # [B, num_anchors, H, W, 5 + num_classes]

    # Sigmoid activation for center, objectness, and class scores
    yolo_output[..., 0:2] = torch.sigmoid(yolo_output[..., 0:2])  # x, y
    yolo_output[..., 4] = torch.sigmoid(yolo_output[..., 4])      # objectness
    yolo_output[..., 5:] = torch.sigmoid(yolo_output[..., 5:])    # class scores

    # Prepare grid for computing bounding box coordinates
    grid_x = torch.arange(grid_w, device=device).repeat(grid_h, 1).view([1, 1, grid_h, grid_w]).float()
    grid_y = torch.arange(grid_h, device=device).repeat(grid_w, 1).t().view([1, 1, grid_h, grid_w]).float()

    # Compute bounding box centers
    pred_boxes = torch.zeros_like(yolo_output[..., 0:4])
    pred_boxes[..., 0] = (yolo_output[..., 0] + grid_x) * self.stride  # x center
    pred_boxes[..., 1] = (yolo_output[..., 1] + grid_y) * self.stride  # y center
    pred_boxes[..., 2] = torch.exp(yolo_output[..., 2]) * self.anchors[:, 0].view(1, num_anchors, 1, 1)  # width
    pred_boxes[..., 3] = torch.exp(yolo_output[..., 3]) * self.anchors[:, 1].view(1, num_anchors, 1, 1)  # height

    # Convert to [x1, y1, x2, y2]
    boxes = torch.zeros_like(yolo_output[..., 0:4])
    boxes[..., 0] = pred_boxes[..., 0] - pred_boxes[..., 2] / 2  # x1
    boxes[..., 1] = pred_boxes[..., 1] - pred_boxes[..., 3] / 2  # y1
    boxes[..., 2] = pred_boxes[..., 0] + pred_boxes[..., 2] / 2  # x2
    boxes[..., 3] = pred_boxes[..., 1] + pred_boxes[..., 3] / 2  # y2

    # Reshape for NMS: [B, num_anchors * H * W, 6 + num_classes]
    boxes = boxes.view(batch_size, -1, 4)
    scores = yolo_output[..., 4].view(batch_size, -1)
    class_scores = yolo_output[..., 5:].view(batch_size, -1, self.num_classes)
    class_conf, class_pred = torch.max(class_scores, dim=-1)  # [B, num_anchors * H * W]

    # Multiply objectness score with class confidence
    scores = scores * class_conf

    # Filter out low scores
    mask = scores > self.conf_threshold
    boxes = boxes[mask]
    scores = scores[mask]
    class_pred = class_pred[mask]

    # Apply Non-Maximum Suppression (NMS)
    keep = torch.ops.torchvision.nms(boxes, scores, self.iou_threshold)

    boxes = boxes[keep]
    scores = scores[keep]
    class_pred = class_pred[keep]

    return YOLODecodeResult(
        boxes=boxes,        
        scores=scores,      
        classes=class_pred 
    )

def encode(self, targets: List[Dict[str, Tensor]]) -> Tensor:
 
    batch_size = len(targets)
    grid_h, grid_w = self.grid_size
    dtype = torch.float32
    device = self.anchors.device

    # Initialize the target tensor
    yolo_targets = torch.zeros(
        (batch_size, self.num_anchors, grid_h, grid_w, 5 + self.num_classes),
        dtype=dtype,
        device=device
    )

    for b, target in enumerate(targets):
        boxes = target['boxes']  # Absolute pixel coordinates [num_objects, 4]
        labels = target['labels']  # Class indices [num_objects]
        confidences = target.get('confidences', torch.ones_like(labels, dtype=dtype))  # [num_objects]
        lengths = target.get('lengths', torch.zeros_like(labels, dtype=dtype))  # [num_objects]

        # Normalize boxes to grid coordinates
        boxes_xy = (boxes[:, 0:2] + boxes[:, 2:4]) / 2.0  # [num_objects, 2]
        boxes_wh = boxes[:, 2:4] - boxes[:, 0:2]  # [num_objects, 2]

        # Convert to relative coordinates with respect to the grid
        scale = torch.tensor([grid_w, grid_h], device=device, dtype=dtype).unsqueeze(0)
        boxes_xy = boxes_xy / self.input_image_size[0] * scale  # [num_objects, 2]
        boxes_wh = boxes_wh / self.input_image_size[0] * scale  # [num_objects, 2]

        # Calculate grid cell indices
        grid_xy = boxes_xy.long()  # [num_objects, 2]
        grid_x = grid_xy[:, 0].clamp(0, grid_w - 1)
        grid_y = grid_xy[:, 1].clamp(0, grid_h - 1)

        # Calculate offsets within grid cells
        x_offset = boxes_xy[:, 0] - grid_x.float()
        y_offset = boxes_xy[:, 1] - grid_y.float()
        box_coords = torch.stack([x_offset, y_offset, boxes_wh[:, 0], boxes_wh[:, 1]], dim=-1)  # [num_objects, 4]

        # Match anchors
        ious = self.calculate_ious(boxes_wh, self.anchors)  # [num_objects, num_anchors]
        best_anchors = ious.argmax(dim=1)  # [num_objects]

        for idx, anchor_idx in enumerate(best_anchors):
            n = anchor_idx
            i = grid_x[idx]
            j = grid_y[idx]
            if j >= grid_h or i >= grid_w:
                continue  # Skip invalid indices

            yolo_targets[b, n, j, i, 0:4] = box_coords[idx]
            yolo_targets[b, n, j, i, 4] = confidences[idx]  # Objectness score
            yolo_targets[b, n, j, i, 5 + labels[idx]] = 1  # One-hot encode class

    return yolo_targets  # [B, num_anchors, grid_h, grid_w, 5 + num_classes]

def calculate_ious(self, boxes_wh: Tensor, anchors_wh: Tensor) -> Tensor:
    
    boxes_area = boxes_wh[:, 0] * boxes_wh[:, 1]  # [num_objects]
    anchors_area = anchors_wh[:, 0] * anchors_wh[:, 1]  # [num_anchors]

    boxes_wh = boxes_wh.unsqueeze(1)  # [num_objects, 1, 2]
    anchors_wh = anchors_wh.unsqueeze(0)  # [1, num_anchors, 2]

    inter_wh = torch.min(boxes_wh, anchors_wh)  # [num_objects, num_anchors, 2]
    inter_area = inter_wh[:, :, 0] * inter_wh[:, :, 1]  # [num_objects, num_anchors]

    union_area = boxes_area.unsqueeze(1) + anchors_area.unsqueeze(0) - inter_area  # [num_objects, num_anchors]
    iou = inter_area / (union_area + 1e-16)  # [num_objects, num_anchors]

    return iou  # [num_objects, num_anchors]
###here's the multilabel_dataset.py which has many issues and missing implements as build_centernet_sample method has missing yolo_targets errors
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,
INPUT_FOLDER_KEY,
TARGET_SAMPLE_WEIGHT,
CENTERNET_TARGET_YOLO,
]

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 and v is not None}
    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)

for k in skip_keys:
    if k == CENTERNET_TARGET_YOLO:
        # Handle CENTERNET_TARGET_YOLO appropriately
        # For example, you might want to stack the tensors
        batch_dict[CENTERNET_TARGET_YOLO] = torch.stack(excluded_items[CENTERNET_TARGET_YOLO])
    else:
        batch_dict[k] = excluded_items[k]

# 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, *args, **kwargs):
    super().__init__(*args, **kwargs)
    # Your initialization code here
    self.stride = kwargs.get('stride', 16)
    self.confidence_mapping = kwargs.get('confidence_mapping', {
        'NONE': 0.0,
        'LOW': 0.33,
        'MEDIUM': 0.66,
        'HIGH': 1.0,
    })
    self.vessel_mapping = kwargs.get('vessel_mapping', {0: 0, 1: 1})
    self.fishing_mapping = kwargs.get('fishing_mapping', {0: 0, 1: 1})
    self.ignore_value = kwargs.get('ignore_value', -1)

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)
    # Ensure 'is_vessel' and 'is_fishing' are the correct data type (strings or numerical labels)
    # Adjust accordingly
    # Keep 'is_vessel' and 'is_fishing' as lists of strings
    is_vessel_list = is_vessel  # List of strings
    is_fishing_list = is_fishing  # List of strings
    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_list) == num_objects,
        len(is_fishing_list) == 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: is_vessel_list,  # Keep as is if they are strings
        TARGET_FISHING_KEY: is_fishing_list,  # Keep as is if they are strings
        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,
        lengths=lengths_np,
        is_vessel=is_vessel_list,
        is_fishing=is_fishing_list,
    )

    # Now, you can safely access `centernet_targets`
    yolo_targets = centernet_targets.yolo_targets

    # If necessary, check if `yolo_targets` is None
    if yolo_targets is None:
        # Initialize yolo_targets appropriately, e.g., with zeros
        # Determine grid_height, grid_width, num_channels based on your model's requirements
        grid_height, grid_width = image.shape[0] // self.stride, image.shape[1] // self.stride
        # Determine num_channels (e.g., objectness + num_classes)
        num_channels = ...  # Replace with your actual number of channels
        yolo_targets = torch.zeros((num_channels, grid_height, grid_width))  # Adjust dimensions as needed

    # Assign the targets to the result dictionary
    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] = yolo_targets

    for key, value in result.items():
        if value is None:
            print(f"Warning: Key '{key}' has a None value in result for index {index}")
    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

The issue you're encountering stems from the fact that yolo_targets are not being generated or populated within your build_centernet_sample method in multilabel_dataset.py. This results in errors when the rest of the pipeline expects yolo_targets to be present.
To resolve this, you need to implement the generation of yolo_targets within the build_centernet_sample method. Since yolo_targets are specific to the YOLO head and decoder, they need to be created based on your dataset annotations (centers, confidences, vessel and fishing labels, lengths, etc.) and aligned with the YOLO model's expectations.
Here's a step-by-step guide to fix the issue:
Modify MultilabelCircleNetTargetsMixin to Include YOLO Target Generation:
Add Necessary Parameters in the __init__ Method:
In your MultilabelCircleNetTargetsMixin class, you'll need to store essential YOLO parameters such as num_classes, anchors, stride, etc.

class MultilabelCircleNetTargetsMixin:
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.stride = kwargs.get('stride', 16)
        self.confidence_mapping = {
            'NONE': 0.0,
            'LOW': 0.33,
            'MEDIUM': 0.66,
            'HIGH': 1.0,
        }
        self.vessel_mapping = {'False': 0, 'True': 1}
        self.fishing_mapping = {'False': 0, 'True': 1}
        self.num_classes = kwargs.get('num_classes', 2)
        anchors = kwargs.get('anchors', [[10,13], [16,30], [33,23], [30,61], [62,45], [59,119]])
        self.anchors = torch.tensor(anchors, dtype=torch.float32)
        self.num_anchors = len(self.anchors)
Implement generate_yolo_targets Method:
Add the generate_yolo_targets Method to the Mixin:
This method will generate the yolo_targets tensor based on your dataset annotations.

class MultilabelCircleNetTargetsMixin:
    # ... existing code ...

    def generate_yolo_targets(self, centers, confidences, is_vessel, is_fishing, lengths, image_size):
        grid_height = image_size[0] // self.stride
        grid_width = image_size[1] // self.stride
        yolo_targets = np.zeros((self.num_anchors, grid_height, grid_width, 5 + self.num_classes), dtype=np.float32)

        for i in range(len(centers)):
            x_center, y_center = centers[i]
            x_center_scaled = x_center / self.stride
            y_center_scaled = y_center / self.stride
            grid_x = int(x_center_scaled)
            grid_y = int(y_center_scaled)

            if grid_x >= grid_width or grid_y >= grid_height:
                continue

            x_offset = x_center_scaled - grid_x
            y_offset = y_center_scaled - grid_y

            confidence_label = confidences[i]
            confidence = self.confidence_mapping.get(confidence_label, 0.0)

            vessel_class = self.vessel_mapping.get(str(is_vessel[i]), 0)

            label_vec = np.zeros(5 + self.num_classes)
            label_vec[0] = x_offset  # x
            label_vec[1] = y_offset  # y
            label_vec[2] = lengths[i] / self.stride  # width
            label_vec[3] = lengths[i] / self.stride  # height
            label_vec[4] = confidence  # objectness score
            label_vec[5 + vessel_class] = 1  # Assign vessel_class

            # Here, you would typically match the object to the best anchor.
            # For simplification, we'll assign the target to all anchors.
            for anchor_idx in range(self.num_anchors):
                yolo_targets[anchor_idx, grid_y, grid_x, :] = label_vec

        return yolo_targets
Note: In a production environment, you should match each ground truth object to the best fitting anchor based on IoU or size similarity. The above code assigns the same target to all anchors for simplicity.
Update the build_centernet_sample Method:
Call generate_yolo_targets Within build_centernet_sample:

class MultilabelCircleNetTargetsMixin:
    # ... existing code ...

    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]:
        image[~np.isfinite(image)] = 0
        box_coder = self.box_coder.create_box_coder_for_image_size(image.shape[:2])

        centers_np = np.array(centers, dtype=np.float32)
        lengths_np = np.array(lengths, dtype=np.float32)
        is_vessel_list = is_vessel
        is_fishing_list = is_fishing
        is_near_shore_np = np.array(is_near_shore, dtype=np.float32)

        num_objects = len(centers_np)
        if not all([
            len(confidences) == num_objects,
            len(lengths_np) == num_objects,
            len(is_vessel_list) == num_objects,
            len(is_fishing_list) == 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: is_vessel_list,
            TARGET_FISHING_KEY: is_fishing_list,
            TARGET_LENGTHS_KEY: torch.from_numpy(lengths_np),
            INPUT_CONFIDENCE_KEY: confidences,
            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 other targets using the box_coder.
        centernet_targets: MultilabelCircleNetEncodeResult = box_coder.encode(
            centers=centers_np,
            confidences=confidences,
            lengths=lengths_np,
            is_vessel=is_vessel_list,
            is_fishing=is_fishing_list,
        )

        # Assign centernet targets to the result.
        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()

        # Now, generate YOLO targets.
        image_size = image.shape[:2]
        yolo_targets = self.generate_yolo_targets(
            centers=centers_np,
            confidences=confidences,
            is_vessel=is_vessel_list,
            is_fishing=is_fishing_list,
            lengths=lengths_np,
            image_size=image_size
        )
        result[CENTERNET_TARGET_YOLO] = torch.from_numpy(yolo_targets).float()

        return result
Ensure Proper Import Statements:
Import Any Necessary Modules or Functions:

import numpy as np
import torch
Note: Since generate_yolo_targets is now a method within your class, you don't need to import it separately.
Adjust the Collate Function if Necessary:
Ensure centernet_collate Handles CENTERNET_TARGET_YOLO:

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,
        INPUT_FOLDER_KEY,
        TARGET_SAMPLE_WEIGHT,
        CENTERNET_TARGET_YOLO,
    ]

    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 and v is not None}
        included_items.append(included)
        for k in skip_keys:
            if k in b:
                excluded_items[k].append(b[k])

    batch_dict = default_collate(included_items)

    for k in skip_keys:
        if k == CENTERNET_TARGET_YOLO:
            batch_dict[CENTERNET_TARGET_YOLO] = torch.stack(excluded_items[CENTERNET_TARGET_YOLO])
        else:
            batch_dict[k] = excluded_items[k]

    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
Validate the Implementation:
Run a Small Test to Ensure yolo_targets are Generated Correctly:

# Example usage:
sample = dataset[0]
print(sample[CENTERNET_TARGET_YOLO].shape)
Ensure the Shapes Match Your Model's Expectations.
Adjust YOLO Loss Computation Accordingly:
In Your Training Loop or Model Forward Pass, Ensure You're Using CENTERNET_TARGET_YOLO for YOLO Loss:

# In your training step:
def training_step(self, batch, batch_idx):
    images = batch[INPUT_IMAGE_KEY]
    yolo_targets = batch[CENTERNET_TARGET_YOLO]
    outputs = self.model(images)

    yolo_output = outputs[CENTERNET_OUTPUT_YOLO]
    loss = self.yolo_loss(yolo_output, yolo_targets)
    # ...
Note: Make sure your model's output keys and loss functions are correctly aligned with your expectations.
Handle Any Missing Implementations or Adjustments:
If yolo_targets Still Raises Issues, Double-Check:
The dimensionality of yolo_targets.
The consistency of class indices and mappings.
The anchoring strategy (assigning ground truth boxes to anchors).
Ensure Consistency Across All Components (Dataset, Model, Loss, Decoder).
Consider Refactoring for Reusability:
If generate_yolo_targets Might Be Useful Elsewhere, You May Move It to a Utility Module.

# In yolo_utils.py
def generate_yolo_targets(...):
    # Implementation

# Then import it where needed
from yolo_utils import generate_yolo_targets
By implementing the generate_yolo_targets method within your dataset class and adjusting the build_centernet_sample method, you ensure that yolo_targets are properly generated and passed through your data pipeline. This should resolve the errors related to missing yolo_targets and allow your integrated YOLO head and decoder to function correctly within your CenterNetUnet framework.
Remember: Thoroughly test each component after changes to catch any mismatches or errors early in the development process.