DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x335684e20>
[DEBUG] Channel 0 shape: (2048, 2048)
[DEBUG] Channel 1 shape: (2048, 2048)
(515.0, 1571.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(1689.0, 1660.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(1767.0, 1840.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(1737.0, 1827.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(421.0, 1517.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(26.0, 1867.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(37.0, 1876.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(122.0, 1465.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(458.0, 1466.0, 9.113393783569336) [<class 'float'>, <class 'float'>, <class 'float'>]
(727.0, 1486.0, 2.026921510696411) [<class 'float'>, <class 'float'>, <class 'float'>]
(465.0, 1657.0, 2.0887291431427) [<class 'float'>, <class 'float'>, <class 'float'>]
(934.0, 853.0, 1.4588253498077393) [<class 'float'>, <class 'float'>, <class 'float'>]
(1048.0, 1818.0, 1.5510203838348389) [<class 'float'>, <class 'float'>, <class 'float'>]
(1491.0, 1586.0, 1.5674598217010498) [<class 'float'>, <class 'float'>, <class 'float'>]
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x3356f2940>
DEBUG:rasterio.env:Starting outermost env
DEBUG:rasterio.env:No GDAL environment exists
DEBUG:rasterio.env:New GDAL environment <rasterio._env.GDALEnv object at 0x3356f2790> created
DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x3356f2790>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x3356f2940>
DEBUG:rasterio._base:Sharing flag: 0
DEBUG:rasterio._base:Nodata success: 1, Nodata value: -32768.000000
DEBUG:rasterio._base:Dataset <open DatasetReader name='/Volumes/Crucial/xview3_dataset/full_dataset/validation/4a97701b4bd81bf7v/VH_dB.tif' mode='r'> is started.
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x3356f2940>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3356f2790> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3356f2790>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x3356f2940>
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x3356f2c40>
DEBUG:rasterio.env:Starting outermost env
DEBUG:rasterio.env:No GDAL environment exists
DEBUG:rasterio.env:New GDAL environment <rasterio._env.GDALEnv object at 0x3356f2790> created
DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x3356f2790>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x3356f2c40>
DEBUG:rasterio._io:Output nodata value read from file: -32768.0
DEBUG:rasterio._io:Output nodata values: [-32768.0]
DEBUG:rasterio._io:all_valid: False
DEBUG:rasterio._io:mask_flags: ([<MaskFlags.nodata: 8>],)
DEBUG:rasterio._io:Jump straight to _read()
DEBUG:rasterio._io:Window: Window(col_off=18432, row_off=4096, width=2048, height=2048)
DEBUG:rasterio._io:IO window xoff=18432.0 yoff=4096.0 width=2048.0 height=2048.0
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x3356f2c40>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3356f2790> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3356f2790>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x3356f2c40>
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x3356f2ac0>
DEBUG:rasterio.env:Starting outermost env
DEBUG:rasterio.env:No GDAL environment exists
DEBUG:rasterio.env:New GDAL environment <rasterio._env.GDALEnv object at 0x3356f2a00> created
DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x3356f2a00>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x3356f2ac0>
DEBUG:rasterio._base:Sharing flag: 0
DEBUG:rasterio._base:Nodata success: 1, Nodata value: -32768.000000
DEBUG:rasterio._base:Dataset <open DatasetReader name='/Volumes/Crucial/xview3_dataset/full_dataset/validation/4a97701b4bd81bf7v/VV_dB.tif' mode='r'> is started.
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x3356f2ac0>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3356f2a00> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3356f2a00>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x3356f2ac0>
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x3356f2c40>
DEBUG:rasterio.env:Starting outermost env
DEBUG:rasterio.env:No GDAL environment exists
DEBUG:rasterio.env:New GDAL environment <rasterio._env.GDALEnv object at 0x3356f2a00> created
DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x3356f2a00>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x3356f2c40>
DEBUG:rasterio._io:Output nodata value read from file: -32768.0
DEBUG:rasterio._io:Output nodata values: [-32768.0]
DEBUG:rasterio._io:all_valid: False
DEBUG:rasterio._io:mask_flags: ([<MaskFlags.nodata: 8>],)
DEBUG:rasterio._io:Jump straight to _read()
DEBUG:rasterio._io:Window: Window(col_off=18432, row_off=4096, width=2048, height=2048)
DEBUG:rasterio._io:IO window xoff=18432.0 yoff=4096.0 width=2048.0 height=2048.0
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x3356f2c40>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3356f2a00> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3356f2a00>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x3356f2c40>
[DEBUG] Channel 0 shape: (2048, 2048)
[DEBUG] Channel 1 shape: (2048, 2048)
(1128.0, 384.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(378.0, 1370.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(1443.0, 305.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(1347.0, 251.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(1153.0, 250.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(361.0, 1177.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(774.0, 831.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(2.0, 1328.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(1311.0, 556.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(38.0, 1461.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(1672.0, 186.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x3356f2730>
DEBUG:rasterio.env:Starting outermost env
DEBUG:rasterio.env:No GDAL environment exists
DEBUG:rasterio.env:New GDAL environment <rasterio._env.GDALEnv object at 0x3356f2880> created
DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x3356f2880>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x3356f2730>
DEBUG:rasterio._base:Sharing flag: 0
DEBUG:rasterio._base:Nodata success: 1, Nodata value: -32768.000000
DEBUG:rasterio._base:Dataset <open DatasetReader name='/Volumes/Crucial/xview3_dataset/full_dataset/validation/4a97701b4bd81bf7v/VH_dB.tif' mode='r'> is started.
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x3356f2730>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3356f2880> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3356f2880>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x3356f2730>
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x335684eb0>
DEBUG:rasterio.env:Starting outermost env
DEBUG:rasterio.env:No GDAL environment exists
DEBUG:rasterio.env:New GDAL environment <rasterio._env.GDALEnv object at 0x3356f2880> created
DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x3356f2880>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x335684eb0>
DEBUG:rasterio._io:Output nodata value read from file: -32768.0
DEBUG:rasterio._io:Output nodata values: [-32768.0]
DEBUG:rasterio._io:all_valid: False
DEBUG:rasterio._io:mask_flags: ([<MaskFlags.nodata: 8>],)
DEBUG:rasterio._io:Jump straight to _read()
DEBUG:rasterio._io:Window: Window(col_off=20480, row_off=4096, width=2048, height=2048)
DEBUG:rasterio._io:IO window xoff=20480.0 yoff=4096.0 width=2048.0 height=2048.0
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x335684eb0>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3356f2880> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3356f2880>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x335684eb0>
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x3356f2ac0>
DEBUG:rasterio.env:Starting outermost env
DEBUG:rasterio.env:No GDAL environment exists
DEBUG:rasterio.env:New GDAL environment <rasterio._env.GDALEnv object at 0x3356f2b50> created
DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x3356f2b50>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x3356f2ac0>
DEBUG:rasterio._base:Sharing flag: 0
DEBUG:rasterio._base:Nodata success: 1, Nodata value: -32768.000000
DEBUG:rasterio._base:Dataset <open DatasetReader name='/Volumes/Crucial/xview3_dataset/full_dataset/validation/4a97701b4bd81bf7v/VV_dB.tif' mode='r'> is started.
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x3356f2ac0>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3356f2b50> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3356f2b50>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x3356f2ac0>
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x335684eb0>
DEBUG:rasterio.env:Starting outermost env
DEBUG:rasterio.env:No GDAL environment exists
DEBUG:rasterio.env:New GDAL environment <rasterio._env.GDALEnv object at 0x3356f2b50> created
DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x3356f2b50>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x335684eb0>
DEBUG:rasterio._io:Output nodata value read from file: -32768.0
DEBUG:rasterio._io:Output nodata values: [-32768.0]
DEBUG:rasterio._io:all_valid: False
DEBUG:rasterio._io:mask_flags: ([<MaskFlags.nodata: 8>],)
DEBUG:rasterio._io:Jump straight to _read()
DEBUG:rasterio._io:Window: Window(col_off=20480, row_off=4096, width=2048, height=2048)
DEBUG:rasterio._io:IO window xoff=20480.0 yoff=4096.0 width=2048.0 height=2048.0
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x335684eb0>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3356f2b50> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3356f2b50>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x335684eb0>
[DEBUG] Channel 0 shape: (2048, 2048)
[DEBUG] Channel 1 shape: (2048, 2048)
(1447.0, 380.0, 12.130756378173828) [<class 'float'>, <class 'float'>, <class 'float'>]
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x3356f2880>
DEBUG:rasterio.env:Starting outermost env
DEBUG:rasterio.env:No GDAL environment exists
DEBUG:rasterio.env:New GDAL environment <rasterio._env.GDALEnv object at 0x3356f2910> created
DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x3356f2910>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x3356f2880>
DEBUG:rasterio._base:Sharing flag: 0
DEBUG:rasterio._base:Nodata success: 1, Nodata value: -32768.000000
DEBUG:rasterio._base:Dataset <open DatasetReader name='/Volumes/Crucial/xview3_dataset/full_dataset/validation/4a97701b4bd81bf7v/VH_dB.tif' mode='r'> is started.
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x3356f2880>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3356f2910> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3356f2910>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x3356f2880>
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x335684eb0>
DEBUG:rasterio.env:Starting outermost env
DEBUG:rasterio.env:No GDAL environment exists
DEBUG:rasterio.env:New GDAL environment <rasterio._env.GDALEnv object at 0x3356f2910> created
DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x3356f2910>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x335684eb0>
DEBUG:rasterio._io:Output nodata value read from file: -32768.0
DEBUG:rasterio._io:Output nodata values: [-32768.0]
DEBUG:rasterio._io:all_valid: False
DEBUG:rasterio._io:mask_flags: ([<MaskFlags.nodata: 8>],)
DEBUG:rasterio._io:Jump straight to _read()
DEBUG:rasterio._io:Window: Window(col_off=22528, row_off=4096, width=2048, height=2048)
DEBUG:rasterio._io:IO window xoff=22528.0 yoff=4096.0 width=2048.0 height=2048.0
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x335684eb0>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3356f2910> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3356f2910>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x335684eb0>
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x3356f2790>
DEBUG:rasterio.env:Starting outermost env
DEBUG:rasterio.env:No GDAL environment exists
DEBUG:rasterio.env:New GDAL environment <rasterio._env.GDALEnv object at 0x3356f2b50> created
DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x3356f2b50>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x3356f2790>
DEBUG:rasterio._base:Sharing flag: 0
DEBUG:rasterio._base:Nodata success: 1, Nodata value: -32768.000000
DEBUG:rasterio._base:Dataset <open DatasetReader name='/Volumes/Crucial/xview3_dataset/full_dataset/validation/4a97701b4bd81bf7v/VV_dB.tif' mode='r'> is started.
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x3356f2790>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3356f2b50> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3356f2b50>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x3356f2790>
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x335684eb0>
DEBUG:rasterio.env:Starting outermost env
DEBUG:rasterio.env:No GDAL environment exists
DEBUG:rasterio.env:New GDAL environment <rasterio._env.GDALEnv object at 0x3356f2b50> created
DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x3356f2b50>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x335684eb0>
DEBUG:rasterio._io:Output nodata value read from file: -32768.0
DEBUG:rasterio._io:Output nodata values: [-32768.0]
DEBUG:rasterio._io:all_valid: False
DEBUG:rasterio._io:mask_flags: ([<MaskFlags.nodata: 8>],)
DEBUG:rasterio._io:Jump straight to _read()
DEBUG:rasterio._io:Window: Window(col_off=22528, row_off=4096, width=2048, height=2048)
DEBUG:rasterio._io:IO window xoff=22528.0 yoff=4096.0 width=2048.0 height=2048.0
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x335684eb0>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3356f2b50> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3356f2b50>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x335684eb0>
[DEBUG] Channel 0 shape: (2048, 2048)
[DEBUG] Channel 1 shape: (2048, 2048)
(482.0, 1681.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(683.0, 1429.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
[DEBUG] Feature 0 shape: torch.Size([4, 24, 1023, 1023])
[DEBUG] Feature 1 shape: torch.Size([4, 40, 512, 512])
[DEBUG] Feature 2 shape: torch.Size([4, 64, 256, 256])
[DEBUG] Feature 3 shape: torch.Size([4, 176, 128, 128])
[DEBUG] Feature 4 shape: torch.Size([4, 512, 64, 64])
[DEBUG] Number of encoder features: 5
[DEBUG] Decoder stage 0
Using encoder feature at index 3
x.shape after concat at stage 0: torch.Size([4, 688, 128, 128])
[DEBUG] Decoder stage 1
Using encoder feature at index 2
x.shape after concat at stage 1: torch.Size([4, 576, 256, 256])
[DEBUG] Decoder stage 2
Using encoder feature at index 1
x.shape after concat at stage 2: torch.Size([4, 424, 512, 512])
[DEBUG] Decoder stage 3
Using encoder feature at index 0
encoder_feature.shape after resizing at stage 3: torch.Size([4, 24, 1024, 1024])
x.shape after concat at stage 3: torch.Size([4, 280, 1024, 1024])
[DEBUG] feature_maps in FPNFuse: type=<class 'list'>
[DEBUG] Valid feature_maps[0] shape: torch.Size([4, 512, 128, 128])
[DEBUG] Valid feature_maps[1] shape: torch.Size([4, 384, 256, 256])
[DEBUG] Valid feature_maps[2] shape: torch.Size([4, 256, 512, 512])
[DEBUG] Valid feature_maps[3] shape: torch.Size([4, 128, 1024, 1024])


#here's multilabel_circle_coder.py

import math
from collections import namedtuple
from typing import Tuple, Optional, Union
from torch.nn import functional as F
import numpy as np
import torch
from torch import Tensor
from typing import List

from .functional import *
from .functional import draw_circle
__all__ = [
    "MultilabelCircleNetCoder",
    "MultilabelCircleNetEncodeResult",
    "MultilabelCircleNetDecodeResult",
]

MultilabelCircleNetEncodeResult = namedtuple(
    "MultilabelCircleNetEncodeResult",
    ["heatmap", "is_fishing", "is_vessel", "lengths", "offset", "yolo_targets"],
)
MultilabelCircleNetDecodeResult = namedtuple(
    "MultilabelCircleNetDecodeResult",
    ["centers", "lengths", "is_fishing", "is_vessel", "scores"],
)


def cast_to_float_if_on_cpu(x: Tensor) -> Tensor:
    if x.device == torch.device("cpu") and x.dtype == torch.float16:
        x = x.float()
    return x


class MultilabelCircleNetCoder:
    output_stride: int
    max_objects: int
    heatmap_encoding: str
    labels_encoding: str
    ignore_value: Optional[int]
    image_height: int
    image_width: int
    image_size: Tuple[int, int]
    fixed_radius: int
    labels_radius: int

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

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

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

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

        # Define mappings for categorical fields
        self.confidence_mapping = {
            'LOW': 0,
            'MEDIUM': 1,
            'HIGH': 2,
            'UNKNOWN': 3  # Additional category if needed
        }
        self.vessel_mapping = {
            'UNKNOWN': 0,
            'TYPE1': 1,
            'TYPE2': 2,
            # Add more vessel types as necessary
        }
        self.fishing_mapping = {
            'UNKNOWN': 0,
            'ACTIVE': 1,
            'INACTIVE': 2,
            # Add more fishing statuses as necessary
        }

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

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

    def encode(
        self,
        centers: np.ndarray,
        confidences: List[str],
        is_vessel: List[str],
        is_fishing: List[str],
        lengths: np.ndarray,
    ) -> MultilabelCircleNetEncodeResult:
        """
        Encode the targets for Centernet.

        Args:
            centers (np.ndarray): [N, 2] array of center coordinates.
            confidences (List[str]): List of confidence levels ('LOW', 'MEDIUM', 'HIGH').
            is_vessel (List[str]): List indicating vessel types ('UNKNOWN', 'TYPE1', 'TYPE2', etc.).
            is_fishing (List[str]): List indicating fishing statuses ('UNKNOWN', 'ACTIVE', 'INACTIVE', etc.).
            lengths (np.ndarray): [N] array of ship lengths in pixels.

        Returns:
            MultilabelCircleNetEncodeResult: Named tuple containing encoded heatmap, vessel map, fishing map, lengths, offset, and YOLO targets.
        """

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

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

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

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

        order = [low_conf_mask, med_conf_mask, hgh_conf_mask]

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

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

        # Create heatmap
        output_height = self.image_height // self.output_stride
        output_width = self.image_width // self.output_stride

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

        # Initialize maps
        is_vessel_map = np.zeros((1, output_height, output_width), dtype=np.float32)
        is_fishing_map = np.zeros((1, output_height, output_width), dtype=np.float32)
        length_map = np.zeros((1, output_height, output_width), dtype=np.float32)
        offset_map = np.zeros((2, output_height, output_width), dtype=np.float32)

        num_objs = len(centers_sorted)

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

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

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

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

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

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

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

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

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

    def generate_yolo_targets(self, centers, confidences, is_vessel, is_fishing, lengths, grid_height, grid_width):
        """
        Generate YOLO-style targets with 6 channels: [x, y, w, h, confidence, class].
        - centers: Ground truth bounding boxes [num_objects, 2].
        - confidences: Ground truth confidence [num_objects].
        - is_vessel: Ground truth vessel types [num_objects].
        - is_fishing: Ground truth fishing statuses [num_objects].
        - lengths: Ground truth lengths [num_objects].
        - grid_height: Height of the output grid.
        - grid_width: Width of the output grid.
        """
        num_objects = len(centers)
        yolo_targets = np.zeros((grid_height, grid_width, 6), dtype=np.float32)

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

            # Convert center coordinates to grid cell format (relative to the grid)
            grid_x = int(center[0] * grid_width)
            grid_y = int(center[1] * grid_height)

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

            # Assign to the correct grid cell
            yolo_targets[grid_y, grid_x, 0:2] = center
            yolo_targets[grid_y, grid_x, 2] = length
            yolo_targets[grid_y, grid_x, 3] = confidence
            yolo_targets[grid_y, grid_x, 4] = vessel_type
            yolo_targets[grid_y, grid_x, 5] = fishing_status

        return yolo_targets

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

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

        length_map = cast_to_float_if_on_cpu(length_map)

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

        if max_objects is None:
            max_objects = self.max_objects

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

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

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

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

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

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

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

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

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

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

#here's also centernet_unet.py

from functools import partial
from typing import List, Dict, Any
from omegaconf import DictConfig
from pytorch_toolbelt.modules import (
    ACT_RELU,
    decoders as D,
    EncoderModule,
    ResidualDeconvolutionUpsample2d,
)
from torch import nn, Tensor
import torch
import timm
from hydra.utils import instantiate
import random
import torch.nn.functional as F
from xview3.centernet.models.custom_encoder import CustomTimmEncoder, DebugUNetDecoder
from .unet_blocks import IRBlock, get_unet_block, SqueezeExcitation
from .necks import FPNNeck
from xview3.nas import ProgressiveNAS
from .heads.yolo_head import YOLOHead
from .transformer_classifier import ConvolutionalClassifier
from xview3.common import FilterResponseNormalization, ResidualBlock, FeaturePyramidNetwork
from ..constants import (
    CENTERNET_OUTPUT_CLASS_MAP,
    CENTERNET_OUTPUT_OBJECTNESS_MAP,
    CENTERNET_OUTPUT_OFFSET,
    CENTERNET_OUTPUT_SIZE,
    CENTERNET_OUTPUT_VESSEL_MAP,
    CENTERNET_OUTPUT_FISHING_MAP,
    CENTERNET_OUTPUT_YOLO,
    CENTERNET_TARGET_YOLO
)
from xview3.centernet.models.heads.decoupled_head import DecoupledHeadGroupNormLateShuffle
from xview3.centernet.models.heads.yolo_head import YOLOHead
from ..losses.circlenet_loss import CircleNetLoss, YOLOLoss
from xview3.centernet.bboxer.multilabel_circle_coder import MultilabelCircleNetCoder

__all__ = ["CenterNetUNetModel"]
from xview3.utils import TorchDevice
device = TorchDevice.choose_torch_device()

class CenterNetUNetModel(nn.Module):
    def __init__(
        self,
        encoder: EncoderModule,
        decoder_features: List[int],
        heads: List[nn.Module],
        box_coder,
        neck_out_channels: int = 256,
        input_upsample_factor: int = 1,
        activation=ACT_RELU,
        num_extra_blocks=0,
        block_type="UnetBlock",
        upsample_type="bilinear",
        extra_block=partial(IRBlock, stride=2),
        num_classes=2,
        anchors=[(10, 13), (16, 30), (33, 23)],
        nas=None,
    ):
        super().__init__()
        
        # Initial Setup
        self.output_stride = getattr(encoder, 'strides', [32])[0]
        self.initial_upsample = (
            nn.UpsamplingBilinear2d(scale_factor=input_upsample_factor)
            if input_upsample_factor != 1 else nn.Identity()
        )
        unet_block = get_unet_block(block_type, activation=activation)
        self.encoder = encoder
        self.extra_encoder_blocks = nn.ModuleList(
            [extra_block(encoder.channels[-1], encoder.channels[-1]) for _ in range(num_extra_blocks)]
        )

        # Upsample Block Selection
        if upsample_type in ['nearest', 'bilinear']:
            upsample_block_cls = partial(
                nn.Upsample,
                scale_factor=2,
                mode=upsample_type,
                align_corners=False
            )
        elif upsample_type == 'rdtsc':
            upsample_block_cls = partial(ResidualDeconvolutionUpsample2d, scale_factor=2)
        elif upsample_type == 'shuffle':
            upsample_block_cls = partial(nn.PixelShuffle, upscale_factor=2)
        else:
            raise ValueError(f"Unknown upsample_type {upsample_type}")

        # Encoder Channels
        if hasattr(encoder, 'channels'):
            encoder_channels = list(encoder.channels)
        else:
            encoder_channels = [32, 24, 40, 112, 320]  # Default channels if not provided
        encoder_channels += [encoder_channels[-1]] * num_extra_blocks

        # Head, NAS, YOLO Head, Classifier Initialization
        self.heads = nn.ModuleList(heads)  # Initialize heads
        self.head_output_channels = getattr(self.heads[0], 'out_channels', 256)
        self.output_channels = self.head_output_channels  # **Added**

        self.nas = nas  # Initialize nas

        self.yolo_head = YOLOHead(num_classes=num_classes, anchors=anchors)
        self.yolo_output_channels = self.yolo_head.num_anchors * (5 + num_classes)  # Adjust as per YOLO implementation

        self.classifier = ConvolutionalClassifier(in_channels=neck_out_channels, num_classes=num_classes)
        self.classifier_output_channels = getattr(self.classifier, 'out_channels', num_classes)

        self.box_coder = box_coder

        # Set neck_output_channels directly without multiplication
        self.neck_output_channels = neck_out_channels

        # Decoder Channels Assignment
        self.decoder_channels = decoder_features

        num_encoder_blocks = len(encoder_channels) - 1 

        if len(self.decoder_channels) < num_encoder_blocks:
            self.decoder_channels.extend(
                [self.decoder_channels[-1]] * (num_encoder_blocks - len(self.decoder_channels))
            )
        elif len(self.decoder_channels) > num_encoder_blocks:
            self.decoder_channels = self.decoder_channels[:num_encoder_blocks]

        self.num_decoded_features = len(encoder_channels)

        # Initialize Decoder
        self.decoder = DebugUNetDecoder(
            encoder_channels=encoder_channels,
            decoder_channels=self.decoder_channels,
            unet_block=unet_block,
            upsample_block=upsample_block_cls,
            center=False,
        )

        # Initialize Neck with correct decoder_channels
        self.neck = FPNNeck(channels=self.decoder_channels, out_channels=neck_out_channels)
        # self.neck_output_channels is already set to neck_out_channels

        # Initialize neck_conv with the correct in_channels
        self.neck_conv = nn.Conv2d(
            in_channels=self.neck_output_channels,  # 256
            out_channels=neck_out_channels,         # 256
            kernel_size=1
        )
        self.batch_norm = FilterResponseNormalization(self.neck_output_channels)

        # Calculate total_input_channels based on the corrected neck_output_channels
        self.total_input_channels = (
            int(self.head_output_channels) +
            int(self.yolo_output_channels) +
            int(self.classifier_output_channels)
        )
        print(f"head_output_channels: {self.head_output_channels}")
        print(f"yolo_output_channels: {self.yolo_output_channels}")
        print(f"classifier_output_channels: {self.classifier_output_channels}")
        print(f"total_input_channels: {self.total_input_channels}")

        # Adjust channels if necessary
        self.yolo_adjust = nn.Conv2d(self.yolo_output_channels, self.head_output_channels, kernel_size=1)
        self.classifier_adjust = nn.Conv2d(self.classifier_output_channels, self.head_output_channels, kernel_size=1)

        #self.fusion = nn.Conv2d(self.head_output_channels * 3, self.head_output_channels, kernel_size=1)
        # Assuming total_input_channels should be 13 (based on your concatenation)
        self.fusion = nn.Conv2d(13, self.head_output_channels, kernel_size=1)

        self.dropout = nn.Dropout(p=0.5)

        # Initialize Loss Functions
        self.circle_net_loss = CircleNetLoss()
        self.yolo_loss = YOLOLoss(num_classes=num_classes, anchors=anchors)
        self.criterion = YOLOLoss(num_classes, anchors)

    def forward(self, x):
        x = self.initial_upsample(x)
        feature_maps = self.encoder(x)
        decoded_features = self.decoder(feature_maps)

        neck_features = self.neck(decoded_features)
        neck_features = self.batch_norm(neck_features)
        neck_features = self.dropout(neck_features)
        neck_features = self.neck_conv(neck_features)

        head_output = self.heads[0](neck_features)
        yolo_output = self.yolo_head(neck_features)
        classification_output = self.classifier(neck_features)

        # Print shapes for debugging
        print(f"yolo_output shape: {yolo_output.shape}")

        # Extract tensors from the head_output dictionary
        objectness_map = head_output[CENTERNET_OUTPUT_OBJECTNESS_MAP]
        vessel_map = head_output[CENTERNET_OUTPUT_VESSEL_MAP]
        fishing_map = head_output[CENTERNET_OUTPUT_FISHING_MAP]
        size_map = head_output[CENTERNET_OUTPUT_SIZE]
        offset_map = head_output[CENTERNET_OUTPUT_OFFSET]
        
        # Adjust indexing for yolo_output
        if len(yolo_output.shape) == 4:
            yolo_output = yolo_output.squeeze(1)  # Remove the extra dimension if it's of size 1

        # Print shapes for debugging
        print(f"objectness_map shape: {objectness_map.shape}")
        print(f"vessel_map shape: {vessel_map.shape}")
        print(f"fishing_map shape: {fishing_map.shape}")
        print(f"size_map shape: {size_map.shape}")
        print(f"offset_map shape: {offset_map.shape}")
        print(f"yolo_output shape: {yolo_output.shape}")
        print(f"classification_output shape: {classification_output.shape}")

        # Determine the target spatial size from the primary outputs
        target_size = objectness_map.size()[-2:]  # (H, W)

        # **Upsample YOLO Output**
        if yolo_output.size()[-2:] != target_size:
            yolo_output = F.interpolate(yolo_output, size=target_size, mode='bilinear', align_corners=False)
            print(f"yolo_output upsampled to: {yolo_output.shape}")

        # **Upsample Classification Output**
        if classification_output.size()[-2:] != target_size:
            classification_output = F.interpolate(classification_output, size=target_size, mode='bilinear', align_corners=False)
            print(f"classification_output upsampled to: {classification_output.shape}")

        # Adjust channels if necessary
        yolo_output = self.yolo_adjust(yolo_output)
        classification_output = self.classifier_adjust(classification_output)

        # **Concatenate all outputs along the channel dimension**
        fused_output = self.fusion(torch.cat([
            objectness_map,
            yolo_output,
            classification_output
        ], dim=1))

        # Ensure yolo_output is included in the outputs dictionary
        outputs = {
            CENTERNET_OUTPUT_OBJECTNESS_MAP: objectness_map,
            CENTERNET_OUTPUT_SIZE: size_map,
            CENTERNET_OUTPUT_OFFSET: offset_map,
            CENTERNET_OUTPUT_VESSEL_MAP: vessel_map,
            CENTERNET_OUTPUT_FISHING_MAP: fishing_map,
            CENTERNET_OUTPUT_YOLO: yolo_output,  # Include YOLO output
        }

        return outputs

    def calculate_loss(self, outputs, targets):
        # Ensure yolo_targets are included in the targets dictionary
        if CENTERNET_TARGET_YOLO not in targets:
            raise KeyError(f"Expected key '{CENTERNET_TARGET_YOLO}' not found in targets.")

        outputs[CENTERNET_TARGET_YOLO] = targets[CENTERNET_TARGET_YOLO]
        circle_net_loss = self.circle_net_loss(**outputs, **targets)

        # Calculate YOLO loss
        yolo_loss = self.yolo_loss(outputs[CENTERNET_OUTPUT_YOLO], outputs[CENTERNET_TARGET_YOLO])

        # Combine losses
        total_loss = circle_net_loss['objectness'] + circle_net_loss['vessel'] + circle_net_loss['fishing'] + circle_net_loss['size'] + yolo_loss

        return total_loss

    @classmethod
    def from_config(cls, **config):
        # Extract and remove 'encoder' from config
        encoder_config = config.pop('encoder')
        num_channels = config.pop('num_channels', 3)
        
        # Extract 'model_name' from encoder_config
        model_name = encoder_config.pop('model_name', 'tf_efficientnet_b5_ns')
        pretrained = encoder_config.pop('pretrained', False)
        
        # Initialize the encoder module
        encoder_module = CustomTimmEncoder(
            model_name=model_name,
            num_channels=num_channels,
            features_only=True,
            pretrained=pretrained
        )
        
        # Adjust encoder input channels
        encoder_module = adjust_encoder_input_channels(encoder_module, num_channels)
        
        # Extract other parameters with pop to remove them from config
        input_upsample_factor = config.pop("input_upsample_factor", 1)
        activation = config.pop('activation', ACT_RELU)
        num_extra_blocks = config.pop('num_extra_blocks', 0)
        num_classes = config.pop("num_classes", 2)
        anchors = config.pop("anchors", [(10, 13), (16, 30), (33, 23)])
        
        # Extract decoder configuration
        decoder_config = config.pop('decoder')
        decoder_features = decoder_config.pop('channels')
        block_type = decoder_config.pop("block_type", "UnetBlock")
        upsample_type = decoder_config.pop("upsample_type", "bilinear")
        
        # Extract neck_out_channels
        neck_out_channels = config.pop('neck_out_channels', 256)
        
        # Extract heads configuration and initialize
        heads_config = config.pop('heads')
        heads = []
        for head_config in heads_config:
            if head_config._target_ == 'xview3.centernet.models.heads.yolo_head.YOLOHead':
                head = instantiate(head_config, num_classes=num_classes, anchors=anchors)
            else:
                head = instantiate(head_config, channels=neck_out_channels)
            heads.append(head)
        
        # Extract box_coder configuration and initialize
        box_coder_config = config.pop('box_coder')
        box_coder = instantiate(
            box_coder_config,
            output_stride=encoder_module.strides[0] // (heads[0].upsample_factor * input_upsample_factor),
        )
        
        # **Define the evaluation function**
        def evaluate_architecture(model, train_loader, val_loader):
            return model.evaluate_architecture(model, train_loader, val_loader)

        # **Initialize ProgressiveNAS with the evaluation function**
        nas_config = config.pop('nas', {})  # Extract nas configuration
        nas = ProgressiveNAS(
            base_model=cls(
                encoder=encoder_module,
                decoder_features=decoder_features,
                heads=heads,
                box_coder=box_coder,
                neck_out_channels=neck_out_channels,
                input_upsample_factor=input_upsample_factor,
                activation=activation,
                block_type=block_type,
                upsample_type=upsample_type,
                num_extra_blocks=num_extra_blocks,
                num_classes=num_classes,  # Pass num_classes
                anchors=anchors,
                **config  # Remaining config without 'neck_out_channels'
            ),
            search_space=nas_config.get('search_space', {}),  # Access search_space from nas_config
            evaluation_function=evaluate_architecture  # Pass the evaluation function
        )

        return cls(
            encoder=encoder_module,
            decoder_features=decoder_features,
            heads=heads,
            box_coder=box_coder,
            nas=nas,  # Pass nas to the class
            neck_out_channels=neck_out_channels,
            input_upsample_factor=input_upsample_factor,
            activation=activation,
            block_type=block_type,
            upsample_type=upsample_type,
            num_extra_blocks=num_extra_blocks,
            num_classes=num_classes,  # Pass num_classes
            anchors=anchors,
            **config  # Remaining config without 'neck_out_channels'
        )

def adjust_encoder_input_channels(encoder_module, in_chans):
    # [Existing code...]
    if hasattr(encoder_module, 'encoder'):
        base_encoder = encoder_module.encoder
    else:
        base_encoder = encoder_module

    if hasattr(base_encoder, 'conv_stem'):
        old_conv = base_encoder.conv_stem
     
        new_conv = nn.Conv2d(
            in_channels=in_chans,
            out_channels=old_conv.out_channels,
            kernel_size=old_conv.kernel_size,
            stride=old_conv.stride,
            padding=old_conv.padding,
            bias=old_conv.bias is not None
        )
        
        with torch.no_grad():
            num_existing_in_channels = min(old_conv.in_channels, in_chans)
            new_conv.weight[:, :num_existing_in_channels, ...] = old_conv.weight[:, :num_existing_in_channels, ...]
            if in_chans > old_conv.in_channels:
                torch.nn.init.kaiming_normal_(new_conv.weight[:, old_conv.in_channels:, ...])

        base_encoder.conv_stem = new_conv
        
        print(f"[DEBUG] After adjustment, encoder.conv_stem.in_channels: {base_encoder.conv_stem.in_channels}")
    else:
        raise NotImplementedError("Adjusting input channels for this encoder is not implemented.")
    return encoder_module

class EnhancedCenterNetUNet(nn.Module):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.backbone = CenterNetUNetModel(in_channels, num_classes)
        self.fpn = FeaturePyramidNetwork([64, 128, 256, 512], 256)
        self.se_blocks = nn.ModuleList([SqueezeExcitation(c) for c in [64, 128, 256, 512]])

    def forward(self, x):
        features = self.backbone(x)
        for i, feature in enumerate(features):
            features[i] = self.se_blocks[i](feature)
        return self.fpn(features)