as a super advanced machine learning software engineer help me fix this error

i use this method and my data are located correctly
export XVIEW3_DIR=/Volumes/Crucial/ 
for example
this folder in error is located
"Volumes/Crucial/data/full/train/2f53eb6658f5d979t"
please check all the models that load the data


the error
python train_multilabel2.py
DEBUG:hydra.core.utils:Setting JobRuntime:name=UNKNOWN_NAME
DEBUG:hydra.core.utils:Setting JobRuntime:name=train_multilabel2
Dataset full_4fold
        ignore_low_confidence_detections True
        ignore_low_confidence_labels     True
        filter_low_confidence_objects    False
Train dataset
        Unknown vessels 7788
        Unknown fishing 31133
Valid dataset
        Unknown vessels 36
        Unknown fishing 3934
Preparing training dataset: 100%|█████████████████████████████████████| 425/425 [00:02<00:00, 192.65it/s]
Preparing validation dataset: 100%|█████████████████████████████████████| 12/12 [00:00<00:00, 533.88it/s]
Available keys in model config: dict_keys(['_target_', 'slug', 'weights', 'pretrained', 'output_stride', 'image_size', 'decoder', 'encoder', 'nas', 'num_channels', 'num_classes', 'anchors', 'input_upsample_factor', 'neck_out_channels', 'box_coder', 'head', 'config'])
Complete model configuration:
 _target_: xview3.centernet.models.CenterNetUNetModel.from_config
slug: b5_unet_s2
weights: IMAGENET1K_V1
pretrained: true
output_stride: 32
image_size:
- 1024
- 1024
decoder:
  _target_: xview3.centernet.models.custom_encoder.DebugUNetDecoder
  encoder_channels:
  - 32
  - 24
  - 40
  - 112
  - 320
  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
  - 2
  - 3
  - 4
nas:
  enabled: true
  search_space:
    num_layers:
    - 2
    - 5
    num_filters:
    - 64
    - 128
    - 256
    kernel_size:
    - 3
    - 5
    activation:
    - relu
    - leaky_relu
    dropout_rate_max: 0.5
  num_trials: 50
num_channels: 2
num_classes: 2
anchors:
- - 10
  - 13
- - 16
  - 30
- - 33
  - 23
input_upsample_factor: 1
neck_out_channels: 256
box_coder:
  _target_: xview3.centernet.MultilabelCircleNetCoder
  image_size:
  - 1024
  - 1024
  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
    input_image_size:
    - 640
    - 640
    stride: 32
config:
  slug: b5_unet_s2
  activation: silu
  num_channels: ${dataset.num_channels}
  num_extra_blocks: 0
  encoder:
    _target_: pytorch_toolbelt.modules.encoders.TimmB5Encoder
    pretrained: true
    activation: silu
    layers:
    - 1
    - 2
    - 3
    - 4
  decoder:
    upsample_type: shuffle
    block_type: IRBlock
    channels:
    - 256
    - 384
    - 512
  head:
    _target_: xview3.centernet.models.heads.DecoupledHeadGroupNormLateShuffle
    classifier_dim: 128
    objectness_dim: 128
    size_dim: 64
    offset_dim: 16
    dropout_rate: 0.1
    activation: silu
    num_groups: 4
    num_blocks: 3
  box_coder:
    _target_: xview3.centernet.MultilabelCircleNetCoder
    image_size:
    - 1024
    - 1024
    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}

[DEBUG] Number of heads instantiated: 2
[DEBUG] Head 0 type: <class 'xview3.centernet.models.heads.decoupled_head.DecoupledHeadGroupNormLateShuffle'>
[DEBUG] Head 1 type: <class 'xview3.centernet.models.heads.yolo_head.YOLOHead'>
[2024-10-04 18:20:05,208][timm.models._builder][INFO] - Loading pretrained weights from Hugging Face hub (timm/efficientnetv2_rw_s.ra2_in1k)
[2024-10-04 18:20:05,662][timm.models._hub][INFO] - [timm/efficientnetv2_rw_s.ra2_in1k] Safe alternative available for 'pytorch_model.bin' (as 'model.safetensors'). Loading weights using safetensors.
[2024-10-04 18:20:05,882][timm.models._builder][WARNING] - Unexpected keys (bn2.bias, bn2.num_batches_tracked, bn2.running_mean, bn2.running_var, bn2.weight, classifier.bias, classifier.weight, conv_head.weight) found while loading pretrained weights. This may be expected if model is being adapted.
[2024-10-04 18:20:06,125][timm.models._builder][INFO] - Loading pretrained weights from Hugging Face hub (timm/efficientnetv2_rw_s.ra2_in1k)
[2024-10-04 18:20:06,551][timm.models._hub][INFO] - [timm/efficientnetv2_rw_s.ra2_in1k] Safe alternative available for 'pytorch_model.bin' (as 'model.safetensors'). Loading weights using safetensors.
[2024-10-04 18:20:06,562][timm.models._builder][INFO] - Converted input conv conv_stem pretrained weights from 3 to 2 channel(s)
[2024-10-04 18:20:06,578][timm.models._builder][WARNING] - Unexpected keys (bn2.bias, bn2.num_batches_tracked, bn2.running_mean, bn2.running_var, bn2.weight, classifier.bias, classifier.weight, conv_head.weight) found while loading pretrained weights. This may be expected if model is being adapted.
[2024-10-04 18:20:06,579][xview3.centernet.models.centernet_unet][INFO] - Adjusted encoder's first conv layer 'initial_conv' to have 2 input channels.
[2024-10-04 18:20:11,977][__main__][INFO] - Total parameters in the model: {'total': 263815591, 'trainable': 263815591, 'encoder': 21657072, 'decoder': 14193536}
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 683, in <module>
    main()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/hydra.py", line 119, in run
    ret = run_job(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 679, in main
    pipeline.train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 441, in train
    self.optimizer, self.scheduler, optimizer_callbacks, self.scaler = self.build_optimizer(self.model, train_loader)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 388, in build_optimizer
    scaler = torch.amp.GradScaler(enabled=optimizer_config.get("fp16", False)) if device != TorchDevice.CPU_DEVICE else None
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/amp/grad_scaler.py", line 132, in __init__
    warnings.warn(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/warnings.py", line 109, in _showwarnmsg
    sw(msg.message, msg.category, msg.filename, msg.lineno,
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 70, in warn_with_traceback
    traceback.print_stack(file=log)
/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/amp/grad_scaler.py:132: UserWarning:

torch.cuda.amp.GradScaler is enabled, but CUDA is not available.  Disabling.

1/100 * Epoch (train):   0% 0/27 [00:00<?, ?it/s]DEBUG:rasterio.session:Could not import boto3, continuing with reduced functionality.
DEBUG:rasterio.session:Could not import boto3, continuing with reduced functionality.
DEBUG:rasterio.session:Could not import boto3, continuing with reduced functionality.
DEBUG:rasterio.session:Could not import boto3, continuing with reduced functionality.
Error executing job with overrides: []
Traceback (most recent call last):
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 679, in main
    pipeline.train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 536, in train
    runner.train(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/runners/runner.py", line 170, in train
    distributed_cmd_run(self.run_experiment, distributed)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/utils/scripts.py", line 176, in distributed_cmd_run
    worker_fn(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/core/runner.py", line 929, in run_experiment
    self._run_event("on_exception")
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/core/runner.py", line 724, in _run_event
    getattr(callback, event)(self)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/callbacks/exception.py", line 34, in on_exception
    raise exception
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/core/runner.py", line 917, in run_experiment
    self._run_stage(stage)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/core/runner.py", line 885, in _run_stage
    self._run_epoch(stage=stage, epoch=self.epoch)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/core/runner.py", line 864, in _run_epoch
    self._run_loader(loader)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/core/runner.py", line 798, in _run_loader
    for i, batch in enumerate(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
FileNotFoundError: Caught FileNotFoundError 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 <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/dataset/random_crop_dataset.py", line 159, in __getitem__
    targets = self.load_image(index, crop_coords=crop_coords)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/dataset/keypoint_dataset.py", line 91, in load_image
    individual_channels = self.read_image_fn(self.scenes[index], channels=self.input_channels, crop_coords=crop_coords)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/dataset/io.py", line 151, in read_multichannel_image
    img = read_tiff(channel_image_fname, crop_coords=crop_coords).astype(np.float32, copy=False)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/dataset/io.py", line 37, in read_tiff
    raise FileNotFoundError(image_fname)
FileNotFoundError: data/full/train/2f53eb6658f5d979t/VH_dB.tif


Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
1/100 * Epoch (train):   0% 0/27 [00:14<?, ?it/s]
(sam3) samerkharboush@Samers-MBP SAR-MPS-OptunaNasFbnYoloTransformers % 

some data loading files

##data_module.py


import os
from collections import namedtuple
from typing import Tuple, Optional, Union, Dict

import numpy as np
import pandas as pd
import torch
from omegaconf import DictConfig
from pytorch_toolbelt.utils import to_numpy

__all__ = [
    "XView3DataModule",
    "ignore_low_confidence_objects",
    "filter_low_confidence_objects",
    "drop_unused_columns",
    "MultilabelTargets",
]

from xview3.constants import PIX_TO_M, IGNORE_LABEL

from .splitters import *


def ignore_low_confidence_objects(df: pd.DataFrame) -> pd.DataFrame:
    """
    Set is_vessel, is_fishing and vessel_length_m fields to nan for low-confidence objects.

    Args:
        df:

    Returns:

    """
    df = df.copy()
    low_confidence_mask = df.confidence == "LOW"
    df.loc[low_confidence_mask, "is_vessel"] = float("nan")
    df.loc[low_confidence_mask, "is_fishing"] = float("nan")
    df.loc[low_confidence_mask, "vessel_length_m"] = float("nan")
    return df


def drop_unused_columns(df):
    return df.drop(columns=["detect_lat", "detect_lon", "top", "left", "bottom", "right"]).copy().reset_index(drop=True)


def filter_low_confidence_objects(df):
    return df[df["confidence"].isin(["HIGH", "MEDIUM"])].copy().reset_index(drop=True)


MultilabelTargets = namedtuple(
    "MultilabelTargets",
    (
        "centers",
        "confidences",
        "is_vessel",
        "is_fishing",
        "lengths",
        "is_near_shore",
        "is_vessel_probs",
        "is_fishing_probs",
        "objectness_probs"
    ),
)


class XView3DataModule:
    def __init__(self, data_dir):
        self.data_dir = data_dir

    @classmethod
    def get_label_targets_from_df(cls, df: pd.DataFrame):
        centers = np.stack([df.detect_scene_column.values, df.detect_scene_row.values], axis=-1)

        labels = cls.encode_labels(df.is_vessel, df.is_fishing)
        lengths = cls.encode_lengths(df.vessel_length_m)

        return centers.reshape((-1, 2)), labels.reshape((-1)), lengths.reshape((-1))

    @classmethod
    def get_multilabel_targets_from_df(cls, df: pd.DataFrame) -> MultilabelTargets:
        centers = np.stack([df.detect_scene_column.values, df.detect_scene_row.values], axis=-1)
        if "confidence" in df:
            confidences = df.confidence.values.copy().reshape((-1))
        else:
            confidences = None

        if "distance_from_shore_km" in df:
            is_near_shore = df.distance_from_shore_km.values <= 2
            is_near_shore = is_near_shore.reshape((-1)).astype(int)
        else:
            is_near_shore = None
        lengths = cls.encode_lengths(df.vessel_length_m)

        is_vessel = df.is_vessel.values.astype(np.float32)
        is_vessel[~np.isfinite(is_vessel)] = IGNORE_LABEL

        is_fishing = df.is_fishing.values.astype(np.float32)
        is_fishing[~np.isfinite(is_fishing)] = IGNORE_LABEL
        # is_fishing[is_vessel == 0] = 0 # If it is non-vessel, then by definition it cannot do fishing

        if "objectness_p" in df:
            objectness_p = df.objectness_p.values
        else:
            objectness_p = None

        if "is_vessel_p" in df:
            is_vessel_p = df.is_vessel_p.values
        else:
            is_vessel_p = None

        if "is_fishing_p" in df:
            is_fishing_p = df.is_fishing_p.values
        else:
            is_fishing_p = None


        return MultilabelTargets(
            centers=centers.reshape((-1, 2)),
            is_vessel=is_vessel.reshape((-1)).astype(int),
            is_fishing=is_fishing.reshape((-1)).astype(int),
            lengths=lengths.reshape((-1)),
            confidences=confidences,
            is_near_shore=is_near_shore,
            is_fishing_probs=is_fishing_p,
            is_vessel_probs=is_vessel_p,
            objectness_probs=objectness_p
        )

    @classmethod
    def encode_lengths(cls, lengths: Union[np.ndarray, pd.Series]) -> np.ndarray:
        """
        Convert length in meters to length in pixels
        :param lengths:
        :return:
        """
        lengths = np.asarray(lengths).astype(dtype=np.float32, copy=True)
        invalid_lengths = ~np.isfinite(lengths)
        lengths[invalid_lengths] = 0
        return lengths / PIX_TO_M

    @classmethod
    def decode_lengths(cls, lengths: Union[np.ndarray, torch.Tensor]) -> np.ndarray:
        return to_numpy(lengths * PIX_TO_M)

    @classmethod
    def encode_labels(cls, is_vessel: Union[np.ndarray, pd.Series], is_fishing: Union[np.ndarray, pd.Series]) -> np.ndarray:
        is_fishing_vec = np.asarray(is_fishing == True).astype(int)
        is_vessel_vec = np.asarray((is_fishing == True) | (is_fishing == False) | (is_vessel == True)).astype(int)
        labels = is_vessel_vec + is_fishing_vec
        return labels

    @classmethod
    def decode_labels(cls, labels: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        is_fishing = labels == 2
        is_vessel = labels != 0
        return is_vessel, is_fishing

    @classmethod
    def iterate_crops(cls, image_size: Tuple[int, int], tile_size: Tuple[int, int], tile_step: Tuple[int, int]):
        image_height, image_width = image_size[:2]
        for y in range(0, image_height - tile_size[0] + 1, tile_step[0]):
            for x in range(0, image_width - tile_size[1] + 1, tile_step[1]):
                yield (y, y + tile_step[0]), (x, x + tile_step[1])

    def train_val_split(
        self, splitter: Union[str, Dict, DictConfig], fold: Optional[int] = None, num_folds: Optional[int] = None
    ) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, str]:
        if isinstance(splitter, str):
            splitter_name = splitter
            splitter_params = {}
        else:
            splitter_name = splitter["name"]
            splitter_params = dict((k, v) for k, v in splitter.items() if k != "name")

        splitter_cls: DatasetSplitter = {
            "precomputed": PrecomputedSplitter,
            "tiny": TinyDatasetSplitter,
            "valid_only": ValidationOnlySplitter,
            "valid_only_with_holdout": ValidationOnlyWithHoldoutSplitter,
            "full": FullDatasetWithHoldoutSplitter,
        }[splitter_name](data_dir=self.data_dir, **splitter_params)
        return splitter_cls.train_test_split(fold=fold, num_folds=num_folds)

    def get_test_scenes(self):
        test_dir = os.path.join(self.data_dir, "test")
        return [os.path.join(test_dir, scene) for scene in os.listdir(test_dir) if os.path.isdir(os.path.join(test_dir, scene))]



##io.py

import os
from typing import List, Dict, Tuple, Optional, Union

import cv2
import numpy as np

__all__ = [
    "stack_multichannel_image",
    "read_tiff",
    "read_multichannel_image",
    "read_tiff_with_scaling",
]

from xview3.constants import NODATA_VV_DB

channel_name_to_file_name = {
    "vv": "VV_dB.tif",  # VV-polarization SAR image (dB), 10-meter pixels
    "vh": "VH_dB.tif",  # VH-polarization SAR image (dB), 10-meter pixels
    "bathymetry": "bathymetry.tif",  # bathymetry and elevation image (meters), 500-meter pixels
    "wind_direction": "owiWindDirection.tif",  # wind direction (degrees clockwise from north), 500-meter pixels
    "wind_quality": "owiWindQuality.tif",  # wind quality mask (0 good, 1 medium, 2 low, 3 poor), 500-meter pixels
    "wind_speed": "owiWindSpeed.tif",  # wind speed (meters per second), 500-meter pixels
    "mask": "owiMask.tif",  # land/ice mask (0 valid, 1 land, 2 ice), 500-meter pixels
}


def read_tiff(image_fname: str, crop_coords: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] = None) -> np.ndarray:
    import rasterio
    from rasterio.windows import Window

    window = None
    if crop_coords is not None:
        (row_start, row_stop), (col_start, col_stop) = crop_coords
        window = Window.from_slices((row_start, row_stop), (col_start, col_stop))

    if not os.path.isfile(image_fname):
        raise FileNotFoundError(image_fname)

    with rasterio.open(image_fname) as f:
        return f.read(1, window=window)


def read_owi_mask_with_scaling(
    image_fname,
    crop_coords: Tuple[Tuple[int, int], Tuple[int, int]],
    s: Union[float, Tuple[float, float]] = 1.0 / 50.0,
) -> np.ndarray:
    """

    :param image:
    :param crop_coords:   (row_start, row_stop), (col_start, col_stop)
    :param s:
    :return:
    """
    if isinstance(s, float):
        sx = sy = s
    else:
        sx, sy = s

    image = read_tiff(image_fname)
    (row_start, row_stop), (col_start, col_stop) = crop_coords

    w = col_stop - col_start
    h = row_stop - row_start
    dsize = w, h

    src = np.array(
        [
            [col_start * sx, row_start * sy],
            [col_stop * sx, row_start * sy],
            [col_start * sx, row_stop * sy],
        ],
        dtype=np.float32,
    )
    # dst = np.array([[0, 0], [w - 1, 0], [0, h - 1]], dtype=np.float32)
    dst = np.array([[0, 0], [w, 0], [0, h]], dtype=np.float32)

    M = cv2.getAffineTransform(np.expand_dims(src, 1), np.expand_dims(dst, 1))

    water = image == 0
    land = (image == 1) | (image == 5)
    ice = (image == 2) | (image == 6)

    patch = cv2.warpAffine(
        np.dstack([water, land, ice]).astype(np.float32),
        M,
        dsize=dsize,
        flags=cv2.INTER_CUBIC,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=0,
    )
    return patch


def read_tiff_with_scaling(
    image_fname: str,
    crop_coords: Tuple[Tuple[int, int], Tuple[int, int]],
    s: Union[float, Tuple[float, float]] = 1.0 / 50.0,
) -> np.ndarray:
    """

    :param image_fname:
    :param crop_coords:   (row_start, row_stop), (col_start, col_stop)
    :param s:
    :return:
    """
    if isinstance(s, float):
        sx = sy = s
    else:
        sx, sy = s

    image = read_tiff(image_fname)
    (row_start, row_stop), (col_start, col_stop) = crop_coords

    w = col_stop - col_start
    h = row_stop - row_start
    dsize = w, h

    src = np.array(
        [
            [col_start * sx, row_start * sy],
            [col_stop * sx, row_start * sy],
            [col_start * sx, row_stop * sy],
        ],
        dtype=np.float32,
    )
    # dst = np.array([[0, 0], [w - 1, 0], [0, h - 1]], dtype=np.float32)
    dst = np.array([[0, 0], [w, 0], [0, h]], dtype=np.float32)

    M = cv2.getAffineTransform(np.expand_dims(src, 1), np.expand_dims(dst, 1))

    patch = cv2.warpAffine(
        image,
        M,
        dsize=dsize,
        flags=cv2.INTER_LANCZOS4,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=NODATA_VV_DB,
    )
    return patch


def read_multichannel_image(scene_id, channels: List[str], crop_coords=None):
    individual_channels = {}
    for channel in channels:
        if channel in {"diff(vv,vh)", "mean(vv,vh)"}:
            continue

        if channel in {"vv", "vh"}:
            channel_image_fname = os.path.join(scene_id, channel_name_to_file_name[channel])
            img = read_tiff(channel_image_fname, crop_coords=crop_coords).astype(np.float32, copy=False)
        elif channel == "sar":
            img = np.dstack(
                [
                    read_tiff(os.path.join(scene_id, channel_name_to_file_name["vh"]), crop_coords=crop_coords).astype(
                        np.float32, copy=False
                    ),
                    read_tiff(os.path.join(scene_id, channel_name_to_file_name["vv"]), crop_coords=crop_coords).astype(
                        np.float32, copy=False
                    ),
                ]
            )
        elif channel == "mask":
            channel_image_fname = os.path.join(scene_id, channel_name_to_file_name[channel])
            img = read_owi_mask_with_scaling(channel_image_fname, crop_coords=crop_coords)
        else:
            channel_image_fname = os.path.join(scene_id, channel_name_to_file_name[channel])
            img = read_tiff_with_scaling(channel_image_fname, crop_coords=crop_coords)
        img[img == NODATA_VV_DB] = float("nan")
        individual_channels[channel] = img

    return individual_channels


def stack_multichannel_image(individual_channels: Dict[str, np.ndarray], input_channels: List[str]):
    channels = []
    for channel_name in input_channels:
        img = individual_channels[channel_name]
        channels.append(img)

    image = np.dstack(channels)
    return image


###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 = confidence_mapping or {0: 0.0, 1: 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,
        grid_height,
        grid_width,
    ):
        num_objects = len(centers)
        yolo_targets = np.zeros(
            (grid_height, grid_width, 6), dtype=np.float32
        )  # Adjust the last dimension as needed
    
        for i in range(num_objects):
            center = centers[i]
            confidence = self.confidence_mapping.get(confidences[i], 0.0)
            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]
    
            # Assuming centers are normalized between 0 and 1
            grid_x = int(center[0] * grid_width)
            grid_y = int(center[1] * grid_height)
    
            # Ensure grid indices are within bounds
            grid_x = min(max(grid_x, 0), grid_width - 1)
            grid_y = min(max(grid_y, 0), 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]:

        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,
            grid_height,
            grid_width
        )

        # 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(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,  # Keep as list of categorical values
            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,  # Change this to a config dictionary
        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,  # Added parameter
        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,
        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_dataset', '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