Skip to content

Basnet

BASNet saliency detection package.

BASNetConfig

Bases: PretrainedConfig

Configuration for BASNet saliency prediction.

Parameters:

Name Type Description Default
id2label Mapping[int | str, str] | None

Public label mapping persisted with the model.

 None
input_size int

Square side length used by the image processor.

 256
rgb_mean Sequence[float]

RGB normalization mean.

 (0.485, 0.456, 0.406)
rgb_std Sequence[float]

RGB normalization standard deviation.

 (0.229, 0.224, 0.225)
conversion_report Mapping[str, str | int | float | bool | list[str]] | None

Conversion metadata persisted in configs.

 None
kwargs str | int | float | bool | None

Extra PretrainedConfig fields.

 {}

Returns:

Type Description

BASNet configuration instance.

Raises:

Type Description
ValueError

If input_size is not positive.

Examples:

>>> config = BASNetConfig(input_size=256)
>>> config.model_type
'basnet'
Source code in models/basnet/src/basnet/configuration_basnet.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
class BASNetConfig(PretrainedConfig):
    """Configuration for BASNet saliency prediction.

    Args:
        id2label: Public label mapping persisted with the model.
        input_size: Square side length used by the image processor.
        rgb_mean: RGB normalization mean.
        rgb_std: RGB normalization standard deviation.
        conversion_report: Conversion metadata persisted in configs.
        kwargs: Extra ``PretrainedConfig`` fields.

    Returns:
        BASNet configuration instance.

    Raises:
        ValueError: If ``input_size`` is not positive.

    Examples:
        >>> config = BASNetConfig(input_size=256)
        >>> config.model_type
        'basnet'
    """

    model_type = "basnet"

    def __init__(
        self,
        *,
        id2label: Mapping[int | str, str] | None = None,
        input_size: int = 256,
        rgb_mean: Sequence[float] = (0.485, 0.456, 0.406),
        rgb_std: Sequence[float] = (0.229, 0.224, 0.225),
        conversion_report: Mapping[str, str | int | float | bool | list[str]]
        | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize BASNet configuration."""
        if input_size <= 0:
            raise ValueError("input_size must be positive")

        raw_id2label = id2label or DEFAULT_ID2LABEL
        normalized_id2label = {int(key): value for key, value in raw_id2label.items()}
        super().__init__(id2label=normalized_id2label, **kwargs)  # ty: ignore[invalid-argument-type]
        self.id2label = normalized_id2label
        self.label2id = {value: key for key, value in self.id2label.items()}
        self.input_size = int(input_size)
        self.rgb_mean = tuple(float(value) for value in rgb_mean)
        self.rgb_std = tuple(float(value) for value in rgb_std)
        self.conversion_report: dict[str, str | int | float | bool | list[str]] = dict(
            conversion_report or {}
        )

__init__ 

__init__(
    *,
    id2label: Mapping[int | str, str] | None = None,
    input_size: int = 256,
    rgb_mean: Sequence[float] = (0.485, 0.456, 0.406),
    rgb_std: Sequence[float] = (0.229, 0.224, 0.225),
    conversion_report: Mapping[
        str, str | int | float | bool | list[str]
    ]
    | None = None,
    **kwargs: str | int | float | bool | None,
) -> None

Initialize BASNet configuration.

Source code in models/basnet/src/basnet/configuration_basnet.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def __init__(
    self,
    *,
    id2label: Mapping[int | str, str] | None = None,
    input_size: int = 256,
    rgb_mean: Sequence[float] = (0.485, 0.456, 0.406),
    rgb_std: Sequence[float] = (0.229, 0.224, 0.225),
    conversion_report: Mapping[str, str | int | float | bool | list[str]]
    | None = None,
    **kwargs: str | int | float | bool | None,
) -> None:
    """Initialize BASNet configuration."""
    if input_size <= 0:
        raise ValueError("input_size must be positive")

    raw_id2label = id2label or DEFAULT_ID2LABEL
    normalized_id2label = {int(key): value for key, value in raw_id2label.items()}
    super().__init__(id2label=normalized_id2label, **kwargs)  # ty: ignore[invalid-argument-type]
    self.id2label = normalized_id2label
    self.label2id = {value: key for key, value in self.id2label.items()}
    self.input_size = int(input_size)
    self.rgb_mean = tuple(float(value) for value in rgb_mean)
    self.rgb_std = tuple(float(value) for value in rgb_std)
    self.conversion_report: dict[str, str | int | float | bool | list[str]] = dict(
        conversion_report or {}
    )

BASNetImageProcessor

Bases: BaseImageProcessor

Prepare BASNet image tensors and image-space saliency maps.

Parameters:

Name Type Description Default
input_size int

Square side length used for model inputs.

 256
rgb_mean Sequence[float]

RGB normalization mean.

 (0.485, 0.456, 0.406)
rgb_std Sequence[float]

RGB normalization standard deviation.

 (0.229, 0.224, 0.225)

Returns:

Type Description

BASNet image processor.

Raises:

Type Description
ValueError

If input_size is not positive.

Examples:

>>> processor = BASNetImageProcessor(input_size=32)
>>> batch = processor.preprocess(Image.new("RGB", (16, 20)))
>>> tuple(batch["pixel_values"].shape)
(1, 3, 32, 32)
Source code in models/basnet/src/basnet/image_processing_basnet.py 
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
class BASNetImageProcessor(BaseImageProcessor):
    """Prepare BASNet image tensors and image-space saliency maps.

    Args:
        input_size: Square side length used for model inputs.
        rgb_mean: RGB normalization mean.
        rgb_std: RGB normalization standard deviation.

    Returns:
        BASNet image processor.

    Raises:
        ValueError: If ``input_size`` is not positive.

    Examples:
        >>> processor = BASNetImageProcessor(input_size=32)
        >>> batch = processor.preprocess(Image.new("RGB", (16, 20)))
        >>> tuple(batch["pixel_values"].shape)
        (1, 3, 32, 32)
    """

    model_input_names = ["pixel_values"]

    def __init__(
        self,
        input_size: int = 256,
        rgb_mean: Sequence[float] = (0.485, 0.456, 0.406),
        rgb_std: Sequence[float] = (0.229, 0.224, 0.225),
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize image processor settings."""
        if input_size <= 0:
            raise ValueError("input_size must be positive")

        super().__init__(**kwargs)
        self.input_size = int(input_size)
        self.rgb_mean = tuple(float(value) for value in rgb_mean)
        self.rgb_std = tuple(float(value) for value in rgb_std)

    @classmethod
    def from_config(cls, config: BASNetConfig) -> "BASNetImageProcessor":
        """Build an image processor from BASNet configuration."""
        return cls(
            input_size=config.input_size,
            rgb_mean=config.rgb_mean,
            rgb_std=config.rgb_std,
        )

    def preprocess(
        self,
        images: ImageInput | Sequence[ImageInput],
        *,
        return_tensors: Literal["pt"] = "pt",
        **kwargs: str | int | float | bool | None,
    ) -> BatchFeature:
        """Preprocess images for BASNet saliency prediction.

        Args:
            images: RGB image or image batch.
            return_tensors: Tensor framework. Only ``pt`` is supported.
            kwargs: Ignored compatibility kwargs.

        Returns:
            Batch feature with ``pixel_values`` and ``image_sizes``.

        Raises:
            ValueError: If ``return_tensors`` is not ``pt``.
            TypeError: If an image input type is unsupported.
        """
        del kwargs
        if return_tensors != "pt":
            raise ValueError("BASNetImageProcessor only supports return_tensors='pt'")

        tensors = []
        sizes = []
        for image in _ensure_pil_batch(images):
            width, height = image.size
            sizes.append((height, width))
            array = resize_basnet_rgb(image.convert("RGB"), self.input_size)
            max_value = float(array.max())
            array = array / (max_value if max_value > 0 else 1.0)
            mean = np.asarray(self.rgb_mean, dtype=array.dtype)
            std = np.asarray(self.rgb_std, dtype=array.dtype)
            tensors.append(torch.from_numpy(((array - mean) / std).transpose(2, 0, 1)))
        return BatchFeature(
            {
                "pixel_values": torch.stack(tensors).float(),
                "image_sizes": torch.tensor(sizes, dtype=torch.long),
            }
        )

    def postprocess_saliency(
        self,
        saliency: Float[torch.Tensor, "height width"]
        | Float[torch.Tensor, "batch height width"],
        *,
        output_size: tuple[int, int] | Sequence[tuple[int, int]],
    ) -> (
        Float[torch.Tensor, "height width"] | Float[torch.Tensor, "batch height width"]
    ):
        """Resize normalized saliency maps through the PNG-space path.

        Args:
            saliency: Normalized saliency map shaped ``(H, W)`` or ``(B, H, W)``.
            output_size: Target ``(height, width)`` or one size per batch row.

        Returns:
            Resized saliency tensor in ``[0, 1]``.

        Raises:
            ValueError: If batch sizes and output sizes do not match.

        Examples:
            >>> processor = BASNetImageProcessor()
            >>> out = processor.postprocess_saliency(torch.zeros(4, 4), output_size=(8, 6))
            >>> tuple(out.shape)
            (8, 6)
        """
        if saliency.ndim == 2:
            if not _is_size(output_size):
                raise ValueError("single saliency map requires one output_size tuple")

            return _resize_saliency_png_space(saliency, output_size)
        if _is_size(output_size):
            sizes = [output_size] * int(saliency.shape[0])
        else:
            sizes = cast(list[tuple[int, int]], list(output_size))
        if len(sizes) != int(saliency.shape[0]):
            raise ValueError("output_size batch length must match saliency batch")

        rows = [
            _resize_saliency_png_space(row, size)
            for row, size in zip(saliency, sizes, strict=True)
        ]
        return torch.stack(rows)

__init__ 

__init__(
    input_size: int = 256,
    rgb_mean: Sequence[float] = (0.485, 0.456, 0.406),
    rgb_std: Sequence[float] = (0.229, 0.224, 0.225),
    **kwargs: str | int | float | bool | None,
) -> None

Initialize image processor settings.

Source code in models/basnet/src/basnet/image_processing_basnet.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def __init__(
    self,
    input_size: int = 256,
    rgb_mean: Sequence[float] = (0.485, 0.456, 0.406),
    rgb_std: Sequence[float] = (0.229, 0.224, 0.225),
    **kwargs: str | int | float | bool | None,
) -> None:
    """Initialize image processor settings."""
    if input_size <= 0:
        raise ValueError("input_size must be positive")

    super().__init__(**kwargs)
    self.input_size = int(input_size)
    self.rgb_mean = tuple(float(value) for value in rgb_mean)
    self.rgb_std = tuple(float(value) for value in rgb_std)

from_config classmethod 

from_config(config: BASNetConfig) -> 'BASNetImageProcessor'

Build an image processor from BASNet configuration.

Source code in models/basnet/src/basnet/image_processing_basnet.py 
58
59
60
61
62
63
64
65
@classmethod
def from_config(cls, config: BASNetConfig) -> "BASNetImageProcessor":
    """Build an image processor from BASNet configuration."""
    return cls(
        input_size=config.input_size,
        rgb_mean=config.rgb_mean,
        rgb_std=config.rgb_std,
    )

preprocess 

preprocess(
    images: ImageInput | Sequence[ImageInput],
    *,
    return_tensors: Literal["pt"] = "pt",
    **kwargs: str | int | float | bool | None,
) -> BatchFeature

Preprocess images for BASNet saliency prediction.

Parameters:

Name Type Description Default
images ImageInput | Sequence[ImageInput]

RGB image or image batch.

 required
return_tensors Literal['pt']

Tensor framework. Only pt is supported.

 'pt'
kwargs str | int | float | bool | None

Ignored compatibility kwargs.

 {}

Returns:

Type Description
BatchFeature

Batch feature with pixel_values and image_sizes.

Raises:

Type Description
ValueError

If return_tensors is not pt.
TypeError

If an image input type is unsupported.
Source code in models/basnet/src/basnet/image_processing_basnet.py 
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def preprocess(
    self,
    images: ImageInput | Sequence[ImageInput],
    *,
    return_tensors: Literal["pt"] = "pt",
    **kwargs: str | int | float | bool | None,
) -> BatchFeature:
    """Preprocess images for BASNet saliency prediction.

    Args:
        images: RGB image or image batch.
        return_tensors: Tensor framework. Only ``pt`` is supported.
        kwargs: Ignored compatibility kwargs.

    Returns:
        Batch feature with ``pixel_values`` and ``image_sizes``.

    Raises:
        ValueError: If ``return_tensors`` is not ``pt``.
        TypeError: If an image input type is unsupported.
    """
    del kwargs
    if return_tensors != "pt":
        raise ValueError("BASNetImageProcessor only supports return_tensors='pt'")

    tensors = []
    sizes = []
    for image in _ensure_pil_batch(images):
        width, height = image.size
        sizes.append((height, width))
        array = resize_basnet_rgb(image.convert("RGB"), self.input_size)
        max_value = float(array.max())
        array = array / (max_value if max_value > 0 else 1.0)
        mean = np.asarray(self.rgb_mean, dtype=array.dtype)
        std = np.asarray(self.rgb_std, dtype=array.dtype)
        tensors.append(torch.from_numpy(((array - mean) / std).transpose(2, 0, 1)))
    return BatchFeature(
        {
            "pixel_values": torch.stack(tensors).float(),
            "image_sizes": torch.tensor(sizes, dtype=torch.long),
        }
    )

postprocess_saliency 

postprocess_saliency(
    saliency: Float[Tensor, "height width"]
    | Float[Tensor, "batch height width"],
    *,
    output_size: tuple[int, int]
    | Sequence[tuple[int, int]],
) -> (
    Float[torch.Tensor, "height width"]
    | Float[torch.Tensor, "batch height width"]
)

Resize normalized saliency maps through the PNG-space path.

Parameters:

Name Type Description Default
saliency Float[Tensor, 'height width'] | Float[Tensor, 'batch height width']

Normalized saliency map shaped (H, W) or (B, H, W).

 required
output_size tuple[int, int] | Sequence[tuple[int, int]]

Target (height, width) or one size per batch row.

 required

Returns:

Type Description
Float[Tensor, 'height width'] | Float[Tensor, 'batch height width']

Resized saliency tensor in [0, 1].

Raises:

Type Description
ValueError

If batch sizes and output sizes do not match.

Examples:

>>> processor = BASNetImageProcessor()
>>> out = processor.postprocess_saliency(torch.zeros(4, 4), output_size=(8, 6))
>>> tuple(out.shape)
(8, 6)
Source code in models/basnet/src/basnet/image_processing_basnet.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def postprocess_saliency(
    self,
    saliency: Float[torch.Tensor, "height width"]
    | Float[torch.Tensor, "batch height width"],
    *,
    output_size: tuple[int, int] | Sequence[tuple[int, int]],
) -> (
    Float[torch.Tensor, "height width"] | Float[torch.Tensor, "batch height width"]
):
    """Resize normalized saliency maps through the PNG-space path.

    Args:
        saliency: Normalized saliency map shaped ``(H, W)`` or ``(B, H, W)``.
        output_size: Target ``(height, width)`` or one size per batch row.

    Returns:
        Resized saliency tensor in ``[0, 1]``.

    Raises:
        ValueError: If batch sizes and output sizes do not match.

    Examples:
        >>> processor = BASNetImageProcessor()
        >>> out = processor.postprocess_saliency(torch.zeros(4, 4), output_size=(8, 6))
        >>> tuple(out.shape)
        (8, 6)
    """
    if saliency.ndim == 2:
        if not _is_size(output_size):
            raise ValueError("single saliency map requires one output_size tuple")

        return _resize_saliency_png_space(saliency, output_size)
    if _is_size(output_size):
        sizes = [output_size] * int(saliency.shape[0])
    else:
        sizes = cast(list[tuple[int, int]], list(output_size))
    if len(sizes) != int(saliency.shape[0]):
        raise ValueError("output_size batch length must match saliency batch")

    rows = [
        _resize_saliency_png_space(row, size)
        for row, size in zip(saliency, sizes, strict=True)
    ]
    return torch.stack(rows)

BASNetModel

Bases: PreTrainedModel

BASNet saliency predictor.

Source code in models/basnet/src/basnet/modeling_basnet.py 
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
class BASNetModel(PreTrainedModel):
    """BASNet saliency predictor."""

    config_class = BASNetConfig
    main_input_name = "pixel_values"
    _tied_weights_keys: list[str] = []

    def __init__(self, config: BASNetConfig) -> None:
        """Initialize the BASNet architecture."""
        super().__init__(config)
        self.all_tied_weights_keys: dict[str, str] = {}
        resnet = models.resnet34(weights=None)
        self.inconv = nn.Conv2d(3, 64, 3, padding=1)
        self.inbn = nn.BatchNorm2d(64)
        self.inrelu = nn.ReLU(inplace=True)

        self.encoder1 = resnet.layer1
        self.encoder2 = resnet.layer2
        self.encoder3 = resnet.layer3
        self.encoder4 = resnet.layer4
        self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True)

        self.resb5_1 = _BasicBlock(512, 512)
        self.resb5_2 = _BasicBlock(512, 512)
        self.resb5_3 = _BasicBlock(512, 512)
        self.pool5 = nn.MaxPool2d(2, 2, ceil_mode=True)

        self.resb6_1 = _BasicBlock(512, 512)
        self.resb6_2 = _BasicBlock(512, 512)
        self.resb6_3 = _BasicBlock(512, 512)

        self.convbg_1 = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
        self.bnbg_1 = nn.BatchNorm2d(512)
        self.relubg_1 = nn.ReLU(inplace=True)
        self.convbg_m = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
        self.bnbg_m = nn.BatchNorm2d(512)
        self.relubg_m = nn.ReLU(inplace=True)
        self.convbg_2 = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
        self.bnbg_2 = nn.BatchNorm2d(512)
        self.relubg_2 = nn.ReLU(inplace=True)

        self.conv6d_1 = nn.Conv2d(1024, 512, 3, padding=1)
        self.bn6d_1 = nn.BatchNorm2d(512)
        self.relu6d_1 = nn.ReLU(inplace=True)
        self.conv6d_m = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
        self.bn6d_m = nn.BatchNorm2d(512)
        self.relu6d_m = nn.ReLU(inplace=True)
        self.conv6d_2 = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
        self.bn6d_2 = nn.BatchNorm2d(512)
        self.relu6d_2 = nn.ReLU(inplace=True)

        self.conv5d_1 = nn.Conv2d(1024, 512, 3, padding=1)
        self.bn5d_1 = nn.BatchNorm2d(512)
        self.relu5d_1 = nn.ReLU(inplace=True)
        self.conv5d_m = nn.Conv2d(512, 512, 3, padding=1)
        self.bn5d_m = nn.BatchNorm2d(512)
        self.relu5d_m = nn.ReLU(inplace=True)
        self.conv5d_2 = nn.Conv2d(512, 512, 3, padding=1)
        self.bn5d_2 = nn.BatchNorm2d(512)
        self.relu5d_2 = nn.ReLU(inplace=True)

        self.conv4d_1 = nn.Conv2d(1024, 512, 3, padding=1)
        self.bn4d_1 = nn.BatchNorm2d(512)
        self.relu4d_1 = nn.ReLU(inplace=True)
        self.conv4d_m = nn.Conv2d(512, 512, 3, padding=1)
        self.bn4d_m = nn.BatchNorm2d(512)
        self.relu4d_m = nn.ReLU(inplace=True)
        self.conv4d_2 = nn.Conv2d(512, 256, 3, padding=1)
        self.bn4d_2 = nn.BatchNorm2d(256)
        self.relu4d_2 = nn.ReLU(inplace=True)

        self.conv3d_1 = nn.Conv2d(512, 256, 3, padding=1)
        self.bn3d_1 = nn.BatchNorm2d(256)
        self.relu3d_1 = nn.ReLU(inplace=True)
        self.conv3d_m = nn.Conv2d(256, 256, 3, padding=1)
        self.bn3d_m = nn.BatchNorm2d(256)
        self.relu3d_m = nn.ReLU(inplace=True)
        self.conv3d_2 = nn.Conv2d(256, 128, 3, padding=1)
        self.bn3d_2 = nn.BatchNorm2d(128)
        self.relu3d_2 = nn.ReLU(inplace=True)

        self.conv2d_1 = nn.Conv2d(256, 128, 3, padding=1)
        self.bn2d_1 = nn.BatchNorm2d(128)
        self.relu2d_1 = nn.ReLU(inplace=True)
        self.conv2d_m = nn.Conv2d(128, 128, 3, padding=1)
        self.bn2d_m = nn.BatchNorm2d(128)
        self.relu2d_m = nn.ReLU(inplace=True)
        self.conv2d_2 = nn.Conv2d(128, 64, 3, padding=1)
        self.bn2d_2 = nn.BatchNorm2d(64)
        self.relu2d_2 = nn.ReLU(inplace=True)

        self.conv1d_1 = nn.Conv2d(128, 64, 3, padding=1)
        self.bn1d_1 = nn.BatchNorm2d(64)
        self.relu1d_1 = nn.ReLU(inplace=True)
        self.conv1d_m = nn.Conv2d(64, 64, 3, padding=1)
        self.bn1d_m = nn.BatchNorm2d(64)
        self.relu1d_m = nn.ReLU(inplace=True)
        self.conv1d_2 = nn.Conv2d(64, 64, 3, padding=1)
        self.bn1d_2 = nn.BatchNorm2d(64)
        self.relu1d_2 = nn.ReLU(inplace=True)

        self.upscore6 = nn.Upsample(
            scale_factor=32, mode="bilinear", align_corners=False
        )
        self.upscore5 = nn.Upsample(
            scale_factor=16, mode="bilinear", align_corners=False
        )
        self.upscore4 = nn.Upsample(
            scale_factor=8, mode="bilinear", align_corners=False
        )
        self.upscore3 = nn.Upsample(
            scale_factor=4, mode="bilinear", align_corners=False
        )
        self.upscore2 = nn.Upsample(
            scale_factor=2, mode="bilinear", align_corners=False
        )

        self.outconvb = nn.Conv2d(512, 1, 3, padding=1)
        self.outconv6 = nn.Conv2d(512, 1, 3, padding=1)
        self.outconv5 = nn.Conv2d(512, 1, 3, padding=1)
        self.outconv4 = nn.Conv2d(256, 1, 3, padding=1)
        self.outconv3 = nn.Conv2d(128, 1, 3, padding=1)
        self.outconv2 = nn.Conv2d(64, 1, 3, padding=1)
        self.outconv1 = nn.Conv2d(64, 1, 3, padding=1)

        self.refunet = _RefUnet(1, 64)

    def _forward_impl(
        self, x: Float[torch.Tensor, "batch channels height width"]
    ) -> tuple[Float[torch.Tensor, "batch channel height width"], ...]:
        hx = self.inrelu(self.inbn(self.inconv(x)))
        h1 = self.encoder1(hx)
        h2 = self.encoder2(h1)
        h3 = self.encoder3(h2)
        h4 = self.encoder4(h3)

        hx = self.pool4(h4)
        hx = self.resb5_1(hx)
        hx = self.resb5_2(hx)
        h5 = self.resb5_3(hx)

        hx = self.pool5(h5)
        hx = self.resb6_1(hx)
        hx = self.resb6_2(hx)
        h6 = self.resb6_3(hx)

        hx = self.relubg_1(self.bnbg_1(self.convbg_1(h6)))
        hx = self.relubg_m(self.bnbg_m(self.convbg_m(hx)))
        hbg = self.relubg_2(self.bnbg_2(self.convbg_2(hx)))

        hx = self.relu6d_1(self.bn6d_1(self.conv6d_1(torch.cat((hbg, h6), 1))))
        hx = self.relu6d_m(self.bn6d_m(self.conv6d_m(hx)))
        hd6 = self.relu6d_2(self.bn5d_2(self.conv6d_2(hx)))
        hx = self.upscore2(hd6)

        hx = self.relu5d_1(self.bn5d_1(self.conv5d_1(torch.cat((hx, h5), 1))))
        hx = self.relu5d_m(self.bn5d_m(self.conv5d_m(hx)))
        hd5 = self.relu5d_2(self.bn5d_2(self.conv5d_2(hx)))
        hx = self.upscore2(hd5)

        hx = self.relu4d_1(self.bn4d_1(self.conv4d_1(torch.cat((hx, h4), 1))))
        hx = self.relu4d_m(self.bn4d_m(self.conv4d_m(hx)))
        hd4 = self.relu4d_2(self.bn4d_2(self.conv4d_2(hx)))
        hx = self.upscore2(hd4)

        hx = self.relu3d_1(self.bn3d_1(self.conv3d_1(torch.cat((hx, h3), 1))))
        hx = self.relu3d_m(self.bn3d_m(self.conv3d_m(hx)))
        hd3 = self.relu3d_2(self.bn3d_2(self.conv3d_2(hx)))
        hx = self.upscore2(hd3)

        hx = self.relu2d_1(self.bn2d_1(self.conv2d_1(torch.cat((hx, h2), 1))))
        hx = self.relu2d_m(self.bn2d_m(self.conv2d_m(hx)))
        hd2 = self.relu2d_2(self.bn2d_2(self.conv2d_2(hx)))
        hx = self.upscore2(hd2)

        hx = self.relu1d_1(self.bn1d_1(self.conv1d_1(torch.cat((hx, h1), 1))))
        hx = self.relu1d_m(self.bn1d_m(self.conv1d_m(hx)))
        hd1 = self.relu1d_2(self.bn1d_2(self.conv1d_2(hx)))

        db = self.upscore6(self.outconvb(hbg))
        d6 = self.upscore6(self.outconv6(hd6))
        d5 = self.upscore5(self.outconv5(hd5))
        d4 = self.upscore4(self.outconv4(hd4))
        d3 = self.upscore3(self.outconv3(hd3))
        d2 = self.upscore2(self.outconv2(hd2))
        d1 = self.outconv1(hd1)
        dout = self.refunet(d1)
        return tuple(torch.sigmoid(row) for row in (dout, d1, d2, d3, d4, d5, d6, db))

    def forward(
        self,
        pixel_values: Float[torch.Tensor, "batch channels height width"],
        return_dict: bool | None = None,
    ) -> BASNetSaliencyOutput | tuple[Shaped[torch.Tensor, "..."], ...]:
        """Predict saliency with the BASNet forward path."""
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )
        outputs = self._forward_impl(pixel_values)
        saliency = normalize_saliency(outputs[0][:, 0, :, :])
        if not return_dict:
            return (saliency, *outputs)
        return BASNetSaliencyOutput(saliency=saliency, side_outputs=outputs)

__init__ 

__init__(config: BASNetConfig) -> None

Initialize the BASNet architecture.

Source code in models/basnet/src/basnet/modeling_basnet.py 
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
def __init__(self, config: BASNetConfig) -> None:
    """Initialize the BASNet architecture."""
    super().__init__(config)
    self.all_tied_weights_keys: dict[str, str] = {}
    resnet = models.resnet34(weights=None)
    self.inconv = nn.Conv2d(3, 64, 3, padding=1)
    self.inbn = nn.BatchNorm2d(64)
    self.inrelu = nn.ReLU(inplace=True)

    self.encoder1 = resnet.layer1
    self.encoder2 = resnet.layer2
    self.encoder3 = resnet.layer3
    self.encoder4 = resnet.layer4
    self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True)

    self.resb5_1 = _BasicBlock(512, 512)
    self.resb5_2 = _BasicBlock(512, 512)
    self.resb5_3 = _BasicBlock(512, 512)
    self.pool5 = nn.MaxPool2d(2, 2, ceil_mode=True)

    self.resb6_1 = _BasicBlock(512, 512)
    self.resb6_2 = _BasicBlock(512, 512)
    self.resb6_3 = _BasicBlock(512, 512)

    self.convbg_1 = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
    self.bnbg_1 = nn.BatchNorm2d(512)
    self.relubg_1 = nn.ReLU(inplace=True)
    self.convbg_m = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
    self.bnbg_m = nn.BatchNorm2d(512)
    self.relubg_m = nn.ReLU(inplace=True)
    self.convbg_2 = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
    self.bnbg_2 = nn.BatchNorm2d(512)
    self.relubg_2 = nn.ReLU(inplace=True)

    self.conv6d_1 = nn.Conv2d(1024, 512, 3, padding=1)
    self.bn6d_1 = nn.BatchNorm2d(512)
    self.relu6d_1 = nn.ReLU(inplace=True)
    self.conv6d_m = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
    self.bn6d_m = nn.BatchNorm2d(512)
    self.relu6d_m = nn.ReLU(inplace=True)
    self.conv6d_2 = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
    self.bn6d_2 = nn.BatchNorm2d(512)
    self.relu6d_2 = nn.ReLU(inplace=True)

    self.conv5d_1 = nn.Conv2d(1024, 512, 3, padding=1)
    self.bn5d_1 = nn.BatchNorm2d(512)
    self.relu5d_1 = nn.ReLU(inplace=True)
    self.conv5d_m = nn.Conv2d(512, 512, 3, padding=1)
    self.bn5d_m = nn.BatchNorm2d(512)
    self.relu5d_m = nn.ReLU(inplace=True)
    self.conv5d_2 = nn.Conv2d(512, 512, 3, padding=1)
    self.bn5d_2 = nn.BatchNorm2d(512)
    self.relu5d_2 = nn.ReLU(inplace=True)

    self.conv4d_1 = nn.Conv2d(1024, 512, 3, padding=1)
    self.bn4d_1 = nn.BatchNorm2d(512)
    self.relu4d_1 = nn.ReLU(inplace=True)
    self.conv4d_m = nn.Conv2d(512, 512, 3, padding=1)
    self.bn4d_m = nn.BatchNorm2d(512)
    self.relu4d_m = nn.ReLU(inplace=True)
    self.conv4d_2 = nn.Conv2d(512, 256, 3, padding=1)
    self.bn4d_2 = nn.BatchNorm2d(256)
    self.relu4d_2 = nn.ReLU(inplace=True)

    self.conv3d_1 = nn.Conv2d(512, 256, 3, padding=1)
    self.bn3d_1 = nn.BatchNorm2d(256)
    self.relu3d_1 = nn.ReLU(inplace=True)
    self.conv3d_m = nn.Conv2d(256, 256, 3, padding=1)
    self.bn3d_m = nn.BatchNorm2d(256)
    self.relu3d_m = nn.ReLU(inplace=True)
    self.conv3d_2 = nn.Conv2d(256, 128, 3, padding=1)
    self.bn3d_2 = nn.BatchNorm2d(128)
    self.relu3d_2 = nn.ReLU(inplace=True)

    self.conv2d_1 = nn.Conv2d(256, 128, 3, padding=1)
    self.bn2d_1 = nn.BatchNorm2d(128)
    self.relu2d_1 = nn.ReLU(inplace=True)
    self.conv2d_m = nn.Conv2d(128, 128, 3, padding=1)
    self.bn2d_m = nn.BatchNorm2d(128)
    self.relu2d_m = nn.ReLU(inplace=True)
    self.conv2d_2 = nn.Conv2d(128, 64, 3, padding=1)
    self.bn2d_2 = nn.BatchNorm2d(64)
    self.relu2d_2 = nn.ReLU(inplace=True)

    self.conv1d_1 = nn.Conv2d(128, 64, 3, padding=1)
    self.bn1d_1 = nn.BatchNorm2d(64)
    self.relu1d_1 = nn.ReLU(inplace=True)
    self.conv1d_m = nn.Conv2d(64, 64, 3, padding=1)
    self.bn1d_m = nn.BatchNorm2d(64)
    self.relu1d_m = nn.ReLU(inplace=True)
    self.conv1d_2 = nn.Conv2d(64, 64, 3, padding=1)
    self.bn1d_2 = nn.BatchNorm2d(64)
    self.relu1d_2 = nn.ReLU(inplace=True)

    self.upscore6 = nn.Upsample(
        scale_factor=32, mode="bilinear", align_corners=False
    )
    self.upscore5 = nn.Upsample(
        scale_factor=16, mode="bilinear", align_corners=False
    )
    self.upscore4 = nn.Upsample(
        scale_factor=8, mode="bilinear", align_corners=False
    )
    self.upscore3 = nn.Upsample(
        scale_factor=4, mode="bilinear", align_corners=False
    )
    self.upscore2 = nn.Upsample(
        scale_factor=2, mode="bilinear", align_corners=False
    )

    self.outconvb = nn.Conv2d(512, 1, 3, padding=1)
    self.outconv6 = nn.Conv2d(512, 1, 3, padding=1)
    self.outconv5 = nn.Conv2d(512, 1, 3, padding=1)
    self.outconv4 = nn.Conv2d(256, 1, 3, padding=1)
    self.outconv3 = nn.Conv2d(128, 1, 3, padding=1)
    self.outconv2 = nn.Conv2d(64, 1, 3, padding=1)
    self.outconv1 = nn.Conv2d(64, 1, 3, padding=1)

    self.refunet = _RefUnet(1, 64)

forward 

forward(
    pixel_values: Float[
        Tensor, "batch channels height width"
    ],
    return_dict: bool | None = None,
) -> (
    BASNetSaliencyOutput
    | tuple[Shaped[torch.Tensor, "..."], ...]
)

Predict saliency with the BASNet forward path.

Source code in models/basnet/src/basnet/modeling_basnet.py 
347
348
349
350
351
352
353
354
355
356
357
358
359
360
def forward(
    self,
    pixel_values: Float[torch.Tensor, "batch channels height width"],
    return_dict: bool | None = None,
) -> BASNetSaliencyOutput | tuple[Shaped[torch.Tensor, "..."], ...]:
    """Predict saliency with the BASNet forward path."""
    return_dict = (
        return_dict if return_dict is not None else self.config.use_return_dict
    )
    outputs = self._forward_impl(pixel_values)
    saliency = normalize_saliency(outputs[0][:, 0, :, :])
    if not return_dict:
        return (saliency, *outputs)
    return BASNetSaliencyOutput(saliency=saliency, side_outputs=outputs)

BASNetSaliencyOutput dataclass

Bases: ModelOutput

Output of BASNetModel.forward.

Source code in models/basnet/src/basnet/modeling_basnet.py 
17
18
19
20
21
22
23
24
@dataclass
class BASNetSaliencyOutput(ModelOutput):
    """Output of ``BASNetModel.forward``."""

    saliency: Float[torch.Tensor, "batch height width"]
    side_outputs: (
        tuple[Float[torch.Tensor, "batch channel height width"], ...] | None
    ) = None

convert_original_checkpoint 

convert_original_checkpoint(
    *,
    checkpoint: Path,
    output_dir: Path,
    config: BASNetConfig,
) -> dict[str, str | int | list[str]]

Convert a raw BASNet checkpoint into a save_pretrained directory.

Parameters:

Name Type Description Default
checkpoint Path

Raw checkpoint path.

 required
output_dir Path

Output model directory.

 required
config BASNetConfig

BASNet config.

 required

Returns:

Type Description
dict[str, str | int | list[str]]

Conversion report dictionary.

Raises:

Type Description
RuntimeError

If converted keys do not strictly match the target model.
Source code in models/basnet/src/basnet/conversion.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def convert_original_checkpoint(
    *,
    checkpoint: Path,
    output_dir: Path,
    config: BASNetConfig,
) -> dict[str, str | int | list[str]]:
    """Convert a raw BASNet checkpoint into a ``save_pretrained`` directory.

    Args:
        checkpoint: Raw checkpoint path.
        output_dir: Output model directory.
        config: BASNet config.

    Returns:
        Conversion report dictionary.

    Raises:
        RuntimeError: If converted keys do not strictly match the target model.
    """
    model = BASNetModel(config)
    state = strip_module_prefix(torch.load(checkpoint, map_location="cpu"))
    missing, unexpected = model.load_state_dict(state, strict=False)
    report: dict[str, str | int | list[str]] = {
        "checkpoint": str(checkpoint),
        "sha256": file_sha256(checkpoint),
        "source_key_count": len(state),
        "missing_keys": [str(key) for key in missing],
        "unexpected_keys": [str(key) for key in unexpected],
    }
    if missing or unexpected:
        raise RuntimeError(json.dumps(report, indent=2, sort_keys=True))

    config.conversion_report = dict[str, str | int | float | bool | list[str]](report)
    model.config = config
    model.save_pretrained(output_dir)
    (output_dir / CONVERSION_REPORT).write_text(
        json.dumps(report, indent=2, sort_keys=True),
        encoding="utf-8",
    )
    return report

normalize_saliency 

normalize_saliency(
    pred: Float[Tensor, "... height width"],
) -> Float[torch.Tensor, "... height width"]

Normalize saliency maps independently over each spatial map.

Source code in models/basnet/src/basnet/modeling_basnet.py 
149
150
151
152
153
154
155
def normalize_saliency(
    pred: Float[torch.Tensor, "... height width"],
) -> Float[torch.Tensor, "... height width"]:
    """Normalize saliency maps independently over each spatial map."""
    min_value = pred.amin(dim=(-2, -1), keepdim=True)
    max_value = pred.amax(dim=(-2, -1), keepdim=True)
    return (pred - min_value) / (max_value - min_value)

configuration_basnet

Configuration for BASNet saliency detection.

BASNetConfig

Bases: PretrainedConfig

Configuration for BASNet saliency prediction.

Parameters:

Name Type Description Default
id2label Mapping[int | str, str] | None

Public label mapping persisted with the model.

 None
input_size int

Square side length used by the image processor.

 256
rgb_mean Sequence[float]

RGB normalization mean.

 (0.485, 0.456, 0.406)
rgb_std Sequence[float]

RGB normalization standard deviation.

 (0.229, 0.224, 0.225)
conversion_report Mapping[str, str | int | float | bool | list[str]] | None

Conversion metadata persisted in configs.

 None
kwargs str | int | float | bool | None

Extra PretrainedConfig fields.

 {}

Returns:

Type Description

BASNet configuration instance.

Raises:

Type Description
ValueError

If input_size is not positive.

Examples:

>>> config = BASNetConfig(input_size=256)
>>> config.model_type
'basnet'
Source code in models/basnet/src/basnet/configuration_basnet.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
class BASNetConfig(PretrainedConfig):
    """Configuration for BASNet saliency prediction.

    Args:
        id2label: Public label mapping persisted with the model.
        input_size: Square side length used by the image processor.
        rgb_mean: RGB normalization mean.
        rgb_std: RGB normalization standard deviation.
        conversion_report: Conversion metadata persisted in configs.
        kwargs: Extra ``PretrainedConfig`` fields.

    Returns:
        BASNet configuration instance.

    Raises:
        ValueError: If ``input_size`` is not positive.

    Examples:
        >>> config = BASNetConfig(input_size=256)
        >>> config.model_type
        'basnet'
    """

    model_type = "basnet"

    def __init__(
        self,
        *,
        id2label: Mapping[int | str, str] | None = None,
        input_size: int = 256,
        rgb_mean: Sequence[float] = (0.485, 0.456, 0.406),
        rgb_std: Sequence[float] = (0.229, 0.224, 0.225),
        conversion_report: Mapping[str, str | int | float | bool | list[str]]
        | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize BASNet configuration."""
        if input_size <= 0:
            raise ValueError("input_size must be positive")

        raw_id2label = id2label or DEFAULT_ID2LABEL
        normalized_id2label = {int(key): value for key, value in raw_id2label.items()}
        super().__init__(id2label=normalized_id2label, **kwargs)  # ty: ignore[invalid-argument-type]
        self.id2label = normalized_id2label
        self.label2id = {value: key for key, value in self.id2label.items()}
        self.input_size = int(input_size)
        self.rgb_mean = tuple(float(value) for value in rgb_mean)
        self.rgb_std = tuple(float(value) for value in rgb_std)
        self.conversion_report: dict[str, str | int | float | bool | list[str]] = dict(
            conversion_report or {}
        )

__init__ 

__init__(
    *,
    id2label: Mapping[int | str, str] | None = None,
    input_size: int = 256,
    rgb_mean: Sequence[float] = (0.485, 0.456, 0.406),
    rgb_std: Sequence[float] = (0.229, 0.224, 0.225),
    conversion_report: Mapping[
        str, str | int | float | bool | list[str]
    ]
    | None = None,
    **kwargs: str | int | float | bool | None,
) -> None

Initialize BASNet configuration.

Source code in models/basnet/src/basnet/configuration_basnet.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def __init__(
    self,
    *,
    id2label: Mapping[int | str, str] | None = None,
    input_size: int = 256,
    rgb_mean: Sequence[float] = (0.485, 0.456, 0.406),
    rgb_std: Sequence[float] = (0.229, 0.224, 0.225),
    conversion_report: Mapping[str, str | int | float | bool | list[str]]
    | None = None,
    **kwargs: str | int | float | bool | None,
) -> None:
    """Initialize BASNet configuration."""
    if input_size <= 0:
        raise ValueError("input_size must be positive")

    raw_id2label = id2label or DEFAULT_ID2LABEL
    normalized_id2label = {int(key): value for key, value in raw_id2label.items()}
    super().__init__(id2label=normalized_id2label, **kwargs)  # ty: ignore[invalid-argument-type]
    self.id2label = normalized_id2label
    self.label2id = {value: key for key, value in self.id2label.items()}
    self.input_size = int(input_size)
    self.rgb_mean = tuple(float(value) for value in rgb_mean)
    self.rgb_std = tuple(float(value) for value in rgb_std)
    self.conversion_report: dict[str, str | int | float | bool | list[str]] = dict(
        conversion_report or {}
    )

conversion

Checkpoint conversion helpers for BASNet.

strip_module_prefix 

strip_module_prefix(
    state_dict: Mapping[str, Shaped[Tensor, "..."]],
) -> dict[str, Shaped[torch.Tensor, "..."]]

Remove DataParallel module. prefixes.

Parameters:

Name Type Description Default
state_dict Mapping[str, Shaped[Tensor, '...']]

Raw PyTorch state dict.

 required

Returns:

Type Description
dict[str, Shaped[Tensor, '...']]

State dict with prefixes removed.

Examples:

>>> strip_module_prefix({"module.a": torch.tensor(1)})["a"].item()
1
Source code in models/basnet/src/basnet/conversion.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def strip_module_prefix(
    state_dict: Mapping[str, Shaped[torch.Tensor, "..."]],
) -> dict[str, Shaped[torch.Tensor, "..."]]:
    """Remove ``DataParallel`` ``module.`` prefixes.

    Args:
        state_dict: Raw PyTorch state dict.

    Returns:
        State dict with prefixes removed.

    Examples:
        >>> strip_module_prefix({"module.a": torch.tensor(1)})["a"].item()
        1
    """
    return {key.removeprefix("module."): value for key, value in state_dict.items()}

file_sha256 

file_sha256(path: Path) -> str

Compute SHA256 for a local file.

Source code in models/basnet/src/basnet/conversion.py 
38
39
40
41
42
43
44
def file_sha256(path: Path) -> str:
    """Compute SHA256 for a local file."""
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()

convert_original_checkpoint 

convert_original_checkpoint(
    *,
    checkpoint: Path,
    output_dir: Path,
    config: BASNetConfig,
) -> dict[str, str | int | list[str]]

Convert a raw BASNet checkpoint into a save_pretrained directory.

Parameters:

Name Type Description Default
checkpoint Path

Raw checkpoint path.

 required
output_dir Path

Output model directory.

 required
config BASNetConfig

BASNet config.

 required

Returns:

Type Description
dict[str, str | int | list[str]]

Conversion report dictionary.

Raises:

Type Description
RuntimeError

If converted keys do not strictly match the target model.
Source code in models/basnet/src/basnet/conversion.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def convert_original_checkpoint(
    *,
    checkpoint: Path,
    output_dir: Path,
    config: BASNetConfig,
) -> dict[str, str | int | list[str]]:
    """Convert a raw BASNet checkpoint into a ``save_pretrained`` directory.

    Args:
        checkpoint: Raw checkpoint path.
        output_dir: Output model directory.
        config: BASNet config.

    Returns:
        Conversion report dictionary.

    Raises:
        RuntimeError: If converted keys do not strictly match the target model.
    """
    model = BASNetModel(config)
    state = strip_module_prefix(torch.load(checkpoint, map_location="cpu"))
    missing, unexpected = model.load_state_dict(state, strict=False)
    report: dict[str, str | int | list[str]] = {
        "checkpoint": str(checkpoint),
        "sha256": file_sha256(checkpoint),
        "source_key_count": len(state),
        "missing_keys": [str(key) for key in missing],
        "unexpected_keys": [str(key) for key in unexpected],
    }
    if missing or unexpected:
        raise RuntimeError(json.dumps(report, indent=2, sort_keys=True))

    config.conversion_report = dict[str, str | int | float | bool | list[str]](report)
    model.config = config
    model.save_pretrained(output_dir)
    (output_dir / CONVERSION_REPORT).write_text(
        json.dumps(report, indent=2, sort_keys=True),
        encoding="utf-8",
    )
    return report

image_processing_basnet

Image processor for BASNet saliency detection.

BASNetImageProcessor

Bases: BaseImageProcessor

Prepare BASNet image tensors and image-space saliency maps.

Parameters:

Name Type Description Default
input_size int

Square side length used for model inputs.

 256
rgb_mean Sequence[float]

RGB normalization mean.

 (0.485, 0.456, 0.406)
rgb_std Sequence[float]

RGB normalization standard deviation.

 (0.229, 0.224, 0.225)

Returns:

Type Description

BASNet image processor.

Raises:

Type Description
ValueError

If input_size is not positive.

Examples:

>>> processor = BASNetImageProcessor(input_size=32)
>>> batch = processor.preprocess(Image.new("RGB", (16, 20)))
>>> tuple(batch["pixel_values"].shape)
(1, 3, 32, 32)
Source code in models/basnet/src/basnet/image_processing_basnet.py 
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
class BASNetImageProcessor(BaseImageProcessor):
    """Prepare BASNet image tensors and image-space saliency maps.

    Args:
        input_size: Square side length used for model inputs.
        rgb_mean: RGB normalization mean.
        rgb_std: RGB normalization standard deviation.

    Returns:
        BASNet image processor.

    Raises:
        ValueError: If ``input_size`` is not positive.

    Examples:
        >>> processor = BASNetImageProcessor(input_size=32)
        >>> batch = processor.preprocess(Image.new("RGB", (16, 20)))
        >>> tuple(batch["pixel_values"].shape)
        (1, 3, 32, 32)
    """

    model_input_names = ["pixel_values"]

    def __init__(
        self,
        input_size: int = 256,
        rgb_mean: Sequence[float] = (0.485, 0.456, 0.406),
        rgb_std: Sequence[float] = (0.229, 0.224, 0.225),
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize image processor settings."""
        if input_size <= 0:
            raise ValueError("input_size must be positive")

        super().__init__(**kwargs)
        self.input_size = int(input_size)
        self.rgb_mean = tuple(float(value) for value in rgb_mean)
        self.rgb_std = tuple(float(value) for value in rgb_std)

    @classmethod
    def from_config(cls, config: BASNetConfig) -> "BASNetImageProcessor":
        """Build an image processor from BASNet configuration."""
        return cls(
            input_size=config.input_size,
            rgb_mean=config.rgb_mean,
            rgb_std=config.rgb_std,
        )

    def preprocess(
        self,
        images: ImageInput | Sequence[ImageInput],
        *,
        return_tensors: Literal["pt"] = "pt",
        **kwargs: str | int | float | bool | None,
    ) -> BatchFeature:
        """Preprocess images for BASNet saliency prediction.

        Args:
            images: RGB image or image batch.
            return_tensors: Tensor framework. Only ``pt`` is supported.
            kwargs: Ignored compatibility kwargs.

        Returns:
            Batch feature with ``pixel_values`` and ``image_sizes``.

        Raises:
            ValueError: If ``return_tensors`` is not ``pt``.
            TypeError: If an image input type is unsupported.
        """
        del kwargs
        if return_tensors != "pt":
            raise ValueError("BASNetImageProcessor only supports return_tensors='pt'")

        tensors = []
        sizes = []
        for image in _ensure_pil_batch(images):
            width, height = image.size
            sizes.append((height, width))
            array = resize_basnet_rgb(image.convert("RGB"), self.input_size)
            max_value = float(array.max())
            array = array / (max_value if max_value > 0 else 1.0)
            mean = np.asarray(self.rgb_mean, dtype=array.dtype)
            std = np.asarray(self.rgb_std, dtype=array.dtype)
            tensors.append(torch.from_numpy(((array - mean) / std).transpose(2, 0, 1)))
        return BatchFeature(
            {
                "pixel_values": torch.stack(tensors).float(),
                "image_sizes": torch.tensor(sizes, dtype=torch.long),
            }
        )

    def postprocess_saliency(
        self,
        saliency: Float[torch.Tensor, "height width"]
        | Float[torch.Tensor, "batch height width"],
        *,
        output_size: tuple[int, int] | Sequence[tuple[int, int]],
    ) -> (
        Float[torch.Tensor, "height width"] | Float[torch.Tensor, "batch height width"]
    ):
        """Resize normalized saliency maps through the PNG-space path.

        Args:
            saliency: Normalized saliency map shaped ``(H, W)`` or ``(B, H, W)``.
            output_size: Target ``(height, width)`` or one size per batch row.

        Returns:
            Resized saliency tensor in ``[0, 1]``.

        Raises:
            ValueError: If batch sizes and output sizes do not match.

        Examples:
            >>> processor = BASNetImageProcessor()
            >>> out = processor.postprocess_saliency(torch.zeros(4, 4), output_size=(8, 6))
            >>> tuple(out.shape)
            (8, 6)
        """
        if saliency.ndim == 2:
            if not _is_size(output_size):
                raise ValueError("single saliency map requires one output_size tuple")

            return _resize_saliency_png_space(saliency, output_size)
        if _is_size(output_size):
            sizes = [output_size] * int(saliency.shape[0])
        else:
            sizes = cast(list[tuple[int, int]], list(output_size))
        if len(sizes) != int(saliency.shape[0]):
            raise ValueError("output_size batch length must match saliency batch")

        rows = [
            _resize_saliency_png_space(row, size)
            for row, size in zip(saliency, sizes, strict=True)
        ]
        return torch.stack(rows)

__init__ 

__init__(
    input_size: int = 256,
    rgb_mean: Sequence[float] = (0.485, 0.456, 0.406),
    rgb_std: Sequence[float] = (0.229, 0.224, 0.225),
    **kwargs: str | int | float | bool | None,
) -> None

Initialize image processor settings.

Source code in models/basnet/src/basnet/image_processing_basnet.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def __init__(
    self,
    input_size: int = 256,
    rgb_mean: Sequence[float] = (0.485, 0.456, 0.406),
    rgb_std: Sequence[float] = (0.229, 0.224, 0.225),
    **kwargs: str | int | float | bool | None,
) -> None:
    """Initialize image processor settings."""
    if input_size <= 0:
        raise ValueError("input_size must be positive")

    super().__init__(**kwargs)
    self.input_size = int(input_size)
    self.rgb_mean = tuple(float(value) for value in rgb_mean)
    self.rgb_std = tuple(float(value) for value in rgb_std)

from_config classmethod 

from_config(config: BASNetConfig) -> 'BASNetImageProcessor'

Build an image processor from BASNet configuration.

Source code in models/basnet/src/basnet/image_processing_basnet.py 
58
59
60
61
62
63
64
65
@classmethod
def from_config(cls, config: BASNetConfig) -> "BASNetImageProcessor":
    """Build an image processor from BASNet configuration."""
    return cls(
        input_size=config.input_size,
        rgb_mean=config.rgb_mean,
        rgb_std=config.rgb_std,
    )

preprocess 

preprocess(
    images: ImageInput | Sequence[ImageInput],
    *,
    return_tensors: Literal["pt"] = "pt",
    **kwargs: str | int | float | bool | None,
) -> BatchFeature

Preprocess images for BASNet saliency prediction.

Parameters:

Name Type Description Default
images ImageInput | Sequence[ImageInput]

RGB image or image batch.

 required
return_tensors Literal['pt']

Tensor framework. Only pt is supported.

 'pt'
kwargs str | int | float | bool | None

Ignored compatibility kwargs.

 {}

Returns:

Type Description
BatchFeature

Batch feature with pixel_values and image_sizes.

Raises:

Type Description
ValueError

If return_tensors is not pt.
TypeError

If an image input type is unsupported.
Source code in models/basnet/src/basnet/image_processing_basnet.py 
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def preprocess(
    self,
    images: ImageInput | Sequence[ImageInput],
    *,
    return_tensors: Literal["pt"] = "pt",
    **kwargs: str | int | float | bool | None,
) -> BatchFeature:
    """Preprocess images for BASNet saliency prediction.

    Args:
        images: RGB image or image batch.
        return_tensors: Tensor framework. Only ``pt`` is supported.
        kwargs: Ignored compatibility kwargs.

    Returns:
        Batch feature with ``pixel_values`` and ``image_sizes``.

    Raises:
        ValueError: If ``return_tensors`` is not ``pt``.
        TypeError: If an image input type is unsupported.
    """
    del kwargs
    if return_tensors != "pt":
        raise ValueError("BASNetImageProcessor only supports return_tensors='pt'")

    tensors = []
    sizes = []
    for image in _ensure_pil_batch(images):
        width, height = image.size
        sizes.append((height, width))
        array = resize_basnet_rgb(image.convert("RGB"), self.input_size)
        max_value = float(array.max())
        array = array / (max_value if max_value > 0 else 1.0)
        mean = np.asarray(self.rgb_mean, dtype=array.dtype)
        std = np.asarray(self.rgb_std, dtype=array.dtype)
        tensors.append(torch.from_numpy(((array - mean) / std).transpose(2, 0, 1)))
    return BatchFeature(
        {
            "pixel_values": torch.stack(tensors).float(),
            "image_sizes": torch.tensor(sizes, dtype=torch.long),
        }
    )

postprocess_saliency 

postprocess_saliency(
    saliency: Float[Tensor, "height width"]
    | Float[Tensor, "batch height width"],
    *,
    output_size: tuple[int, int]
    | Sequence[tuple[int, int]],
) -> (
    Float[torch.Tensor, "height width"]
    | Float[torch.Tensor, "batch height width"]
)

Resize normalized saliency maps through the PNG-space path.

Parameters:

Name Type Description Default
saliency Float[Tensor, 'height width'] | Float[Tensor, 'batch height width']

Normalized saliency map shaped (H, W) or (B, H, W).

 required
output_size tuple[int, int] | Sequence[tuple[int, int]]

Target (height, width) or one size per batch row.

 required

Returns:

Type Description
Float[Tensor, 'height width'] | Float[Tensor, 'batch height width']

Resized saliency tensor in [0, 1].

Raises:

Type Description
ValueError

If batch sizes and output sizes do not match.

Examples:

>>> processor = BASNetImageProcessor()
>>> out = processor.postprocess_saliency(torch.zeros(4, 4), output_size=(8, 6))
>>> tuple(out.shape)
(8, 6)
Source code in models/basnet/src/basnet/image_processing_basnet.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def postprocess_saliency(
    self,
    saliency: Float[torch.Tensor, "height width"]
    | Float[torch.Tensor, "batch height width"],
    *,
    output_size: tuple[int, int] | Sequence[tuple[int, int]],
) -> (
    Float[torch.Tensor, "height width"] | Float[torch.Tensor, "batch height width"]
):
    """Resize normalized saliency maps through the PNG-space path.

    Args:
        saliency: Normalized saliency map shaped ``(H, W)`` or ``(B, H, W)``.
        output_size: Target ``(height, width)`` or one size per batch row.

    Returns:
        Resized saliency tensor in ``[0, 1]``.

    Raises:
        ValueError: If batch sizes and output sizes do not match.

    Examples:
        >>> processor = BASNetImageProcessor()
        >>> out = processor.postprocess_saliency(torch.zeros(4, 4), output_size=(8, 6))
        >>> tuple(out.shape)
        (8, 6)
    """
    if saliency.ndim == 2:
        if not _is_size(output_size):
            raise ValueError("single saliency map requires one output_size tuple")

        return _resize_saliency_png_space(saliency, output_size)
    if _is_size(output_size):
        sizes = [output_size] * int(saliency.shape[0])
    else:
        sizes = cast(list[tuple[int, int]], list(output_size))
    if len(sizes) != int(saliency.shape[0]):
        raise ValueError("output_size batch length must match saliency batch")

    rows = [
        _resize_saliency_png_space(row, size)
        for row, size in zip(saliency, sizes, strict=True)
    ]
    return torch.stack(rows)

resize_basnet_rgb 

resize_basnet_rgb(
    image: Image, input_size: int = 256
) -> Float[np.ndarray, "height width channels"]

Resize an RGB image to the BASNet square input size.

Source code in models/basnet/src/basnet/image_processing_basnet.py 
156
157
158
159
160
161
162
163
164
165
166
167
def resize_basnet_rgb(
    image: Image.Image,
    input_size: int = 256,
) -> Float[np.ndarray, "height width channels"]:
    """Resize an RGB image to the BASNet square input size."""
    raw = np.asarray(image)
    try:
        from skimage import transform  # type: ignore[import-not-found]
    except ModuleNotFoundError:
        resized = image.resize((input_size, input_size), Image.Resampling.BILINEAR)
        return np.asarray(resized, dtype=np.float32)
    return transform.resize(raw, (input_size, input_size), mode="constant")

modeling_basnet

BASNet saliency model.

BASNetSaliencyOutput dataclass

Bases: ModelOutput

Output of BASNetModel.forward.

Source code in models/basnet/src/basnet/modeling_basnet.py 
17
18
19
20
21
22
23
24
@dataclass
class BASNetSaliencyOutput(ModelOutput):
    """Output of ``BASNetModel.forward``."""

    saliency: Float[torch.Tensor, "batch height width"]
    side_outputs: (
        tuple[Float[torch.Tensor, "batch channel height width"], ...] | None
    ) = None

BASNetModel

Bases: PreTrainedModel

BASNet saliency predictor.

Source code in models/basnet/src/basnet/modeling_basnet.py 
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
class BASNetModel(PreTrainedModel):
    """BASNet saliency predictor."""

    config_class = BASNetConfig
    main_input_name = "pixel_values"
    _tied_weights_keys: list[str] = []

    def __init__(self, config: BASNetConfig) -> None:
        """Initialize the BASNet architecture."""
        super().__init__(config)
        self.all_tied_weights_keys: dict[str, str] = {}
        resnet = models.resnet34(weights=None)
        self.inconv = nn.Conv2d(3, 64, 3, padding=1)
        self.inbn = nn.BatchNorm2d(64)
        self.inrelu = nn.ReLU(inplace=True)

        self.encoder1 = resnet.layer1
        self.encoder2 = resnet.layer2
        self.encoder3 = resnet.layer3
        self.encoder4 = resnet.layer4
        self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True)

        self.resb5_1 = _BasicBlock(512, 512)
        self.resb5_2 = _BasicBlock(512, 512)
        self.resb5_3 = _BasicBlock(512, 512)
        self.pool5 = nn.MaxPool2d(2, 2, ceil_mode=True)

        self.resb6_1 = _BasicBlock(512, 512)
        self.resb6_2 = _BasicBlock(512, 512)
        self.resb6_3 = _BasicBlock(512, 512)

        self.convbg_1 = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
        self.bnbg_1 = nn.BatchNorm2d(512)
        self.relubg_1 = nn.ReLU(inplace=True)
        self.convbg_m = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
        self.bnbg_m = nn.BatchNorm2d(512)
        self.relubg_m = nn.ReLU(inplace=True)
        self.convbg_2 = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
        self.bnbg_2 = nn.BatchNorm2d(512)
        self.relubg_2 = nn.ReLU(inplace=True)

        self.conv6d_1 = nn.Conv2d(1024, 512, 3, padding=1)
        self.bn6d_1 = nn.BatchNorm2d(512)
        self.relu6d_1 = nn.ReLU(inplace=True)
        self.conv6d_m = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
        self.bn6d_m = nn.BatchNorm2d(512)
        self.relu6d_m = nn.ReLU(inplace=True)
        self.conv6d_2 = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
        self.bn6d_2 = nn.BatchNorm2d(512)
        self.relu6d_2 = nn.ReLU(inplace=True)

        self.conv5d_1 = nn.Conv2d(1024, 512, 3, padding=1)
        self.bn5d_1 = nn.BatchNorm2d(512)
        self.relu5d_1 = nn.ReLU(inplace=True)
        self.conv5d_m = nn.Conv2d(512, 512, 3, padding=1)
        self.bn5d_m = nn.BatchNorm2d(512)
        self.relu5d_m = nn.ReLU(inplace=True)
        self.conv5d_2 = nn.Conv2d(512, 512, 3, padding=1)
        self.bn5d_2 = nn.BatchNorm2d(512)
        self.relu5d_2 = nn.ReLU(inplace=True)

        self.conv4d_1 = nn.Conv2d(1024, 512, 3, padding=1)
        self.bn4d_1 = nn.BatchNorm2d(512)
        self.relu4d_1 = nn.ReLU(inplace=True)
        self.conv4d_m = nn.Conv2d(512, 512, 3, padding=1)
        self.bn4d_m = nn.BatchNorm2d(512)
        self.relu4d_m = nn.ReLU(inplace=True)
        self.conv4d_2 = nn.Conv2d(512, 256, 3, padding=1)
        self.bn4d_2 = nn.BatchNorm2d(256)
        self.relu4d_2 = nn.ReLU(inplace=True)

        self.conv3d_1 = nn.Conv2d(512, 256, 3, padding=1)
        self.bn3d_1 = nn.BatchNorm2d(256)
        self.relu3d_1 = nn.ReLU(inplace=True)
        self.conv3d_m = nn.Conv2d(256, 256, 3, padding=1)
        self.bn3d_m = nn.BatchNorm2d(256)
        self.relu3d_m = nn.ReLU(inplace=True)
        self.conv3d_2 = nn.Conv2d(256, 128, 3, padding=1)
        self.bn3d_2 = nn.BatchNorm2d(128)
        self.relu3d_2 = nn.ReLU(inplace=True)

        self.conv2d_1 = nn.Conv2d(256, 128, 3, padding=1)
        self.bn2d_1 = nn.BatchNorm2d(128)
        self.relu2d_1 = nn.ReLU(inplace=True)
        self.conv2d_m = nn.Conv2d(128, 128, 3, padding=1)
        self.bn2d_m = nn.BatchNorm2d(128)
        self.relu2d_m = nn.ReLU(inplace=True)
        self.conv2d_2 = nn.Conv2d(128, 64, 3, padding=1)
        self.bn2d_2 = nn.BatchNorm2d(64)
        self.relu2d_2 = nn.ReLU(inplace=True)

        self.conv1d_1 = nn.Conv2d(128, 64, 3, padding=1)
        self.bn1d_1 = nn.BatchNorm2d(64)
        self.relu1d_1 = nn.ReLU(inplace=True)
        self.conv1d_m = nn.Conv2d(64, 64, 3, padding=1)
        self.bn1d_m = nn.BatchNorm2d(64)
        self.relu1d_m = nn.ReLU(inplace=True)
        self.conv1d_2 = nn.Conv2d(64, 64, 3, padding=1)
        self.bn1d_2 = nn.BatchNorm2d(64)
        self.relu1d_2 = nn.ReLU(inplace=True)

        self.upscore6 = nn.Upsample(
            scale_factor=32, mode="bilinear", align_corners=False
        )
        self.upscore5 = nn.Upsample(
            scale_factor=16, mode="bilinear", align_corners=False
        )
        self.upscore4 = nn.Upsample(
            scale_factor=8, mode="bilinear", align_corners=False
        )
        self.upscore3 = nn.Upsample(
            scale_factor=4, mode="bilinear", align_corners=False
        )
        self.upscore2 = nn.Upsample(
            scale_factor=2, mode="bilinear", align_corners=False
        )

        self.outconvb = nn.Conv2d(512, 1, 3, padding=1)
        self.outconv6 = nn.Conv2d(512, 1, 3, padding=1)
        self.outconv5 = nn.Conv2d(512, 1, 3, padding=1)
        self.outconv4 = nn.Conv2d(256, 1, 3, padding=1)
        self.outconv3 = nn.Conv2d(128, 1, 3, padding=1)
        self.outconv2 = nn.Conv2d(64, 1, 3, padding=1)
        self.outconv1 = nn.Conv2d(64, 1, 3, padding=1)

        self.refunet = _RefUnet(1, 64)

    def _forward_impl(
        self, x: Float[torch.Tensor, "batch channels height width"]
    ) -> tuple[Float[torch.Tensor, "batch channel height width"], ...]:
        hx = self.inrelu(self.inbn(self.inconv(x)))
        h1 = self.encoder1(hx)
        h2 = self.encoder2(h1)
        h3 = self.encoder3(h2)
        h4 = self.encoder4(h3)

        hx = self.pool4(h4)
        hx = self.resb5_1(hx)
        hx = self.resb5_2(hx)
        h5 = self.resb5_3(hx)

        hx = self.pool5(h5)
        hx = self.resb6_1(hx)
        hx = self.resb6_2(hx)
        h6 = self.resb6_3(hx)

        hx = self.relubg_1(self.bnbg_1(self.convbg_1(h6)))
        hx = self.relubg_m(self.bnbg_m(self.convbg_m(hx)))
        hbg = self.relubg_2(self.bnbg_2(self.convbg_2(hx)))

        hx = self.relu6d_1(self.bn6d_1(self.conv6d_1(torch.cat((hbg, h6), 1))))
        hx = self.relu6d_m(self.bn6d_m(self.conv6d_m(hx)))
        hd6 = self.relu6d_2(self.bn5d_2(self.conv6d_2(hx)))
        hx = self.upscore2(hd6)

        hx = self.relu5d_1(self.bn5d_1(self.conv5d_1(torch.cat((hx, h5), 1))))
        hx = self.relu5d_m(self.bn5d_m(self.conv5d_m(hx)))
        hd5 = self.relu5d_2(self.bn5d_2(self.conv5d_2(hx)))
        hx = self.upscore2(hd5)

        hx = self.relu4d_1(self.bn4d_1(self.conv4d_1(torch.cat((hx, h4), 1))))
        hx = self.relu4d_m(self.bn4d_m(self.conv4d_m(hx)))
        hd4 = self.relu4d_2(self.bn4d_2(self.conv4d_2(hx)))
        hx = self.upscore2(hd4)

        hx = self.relu3d_1(self.bn3d_1(self.conv3d_1(torch.cat((hx, h3), 1))))
        hx = self.relu3d_m(self.bn3d_m(self.conv3d_m(hx)))
        hd3 = self.relu3d_2(self.bn3d_2(self.conv3d_2(hx)))
        hx = self.upscore2(hd3)

        hx = self.relu2d_1(self.bn2d_1(self.conv2d_1(torch.cat((hx, h2), 1))))
        hx = self.relu2d_m(self.bn2d_m(self.conv2d_m(hx)))
        hd2 = self.relu2d_2(self.bn2d_2(self.conv2d_2(hx)))
        hx = self.upscore2(hd2)

        hx = self.relu1d_1(self.bn1d_1(self.conv1d_1(torch.cat((hx, h1), 1))))
        hx = self.relu1d_m(self.bn1d_m(self.conv1d_m(hx)))
        hd1 = self.relu1d_2(self.bn1d_2(self.conv1d_2(hx)))

        db = self.upscore6(self.outconvb(hbg))
        d6 = self.upscore6(self.outconv6(hd6))
        d5 = self.upscore5(self.outconv5(hd5))
        d4 = self.upscore4(self.outconv4(hd4))
        d3 = self.upscore3(self.outconv3(hd3))
        d2 = self.upscore2(self.outconv2(hd2))
        d1 = self.outconv1(hd1)
        dout = self.refunet(d1)
        return tuple(torch.sigmoid(row) for row in (dout, d1, d2, d3, d4, d5, d6, db))

    def forward(
        self,
        pixel_values: Float[torch.Tensor, "batch channels height width"],
        return_dict: bool | None = None,
    ) -> BASNetSaliencyOutput | tuple[Shaped[torch.Tensor, "..."], ...]:
        """Predict saliency with the BASNet forward path."""
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )
        outputs = self._forward_impl(pixel_values)
        saliency = normalize_saliency(outputs[0][:, 0, :, :])
        if not return_dict:
            return (saliency, *outputs)
        return BASNetSaliencyOutput(saliency=saliency, side_outputs=outputs)

__init__ 

__init__(config: BASNetConfig) -> None

Initialize the BASNet architecture.

Source code in models/basnet/src/basnet/modeling_basnet.py 
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
def __init__(self, config: BASNetConfig) -> None:
    """Initialize the BASNet architecture."""
    super().__init__(config)
    self.all_tied_weights_keys: dict[str, str] = {}
    resnet = models.resnet34(weights=None)
    self.inconv = nn.Conv2d(3, 64, 3, padding=1)
    self.inbn = nn.BatchNorm2d(64)
    self.inrelu = nn.ReLU(inplace=True)

    self.encoder1 = resnet.layer1
    self.encoder2 = resnet.layer2
    self.encoder3 = resnet.layer3
    self.encoder4 = resnet.layer4
    self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True)

    self.resb5_1 = _BasicBlock(512, 512)
    self.resb5_2 = _BasicBlock(512, 512)
    self.resb5_3 = _BasicBlock(512, 512)
    self.pool5 = nn.MaxPool2d(2, 2, ceil_mode=True)

    self.resb6_1 = _BasicBlock(512, 512)
    self.resb6_2 = _BasicBlock(512, 512)
    self.resb6_3 = _BasicBlock(512, 512)

    self.convbg_1 = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
    self.bnbg_1 = nn.BatchNorm2d(512)
    self.relubg_1 = nn.ReLU(inplace=True)
    self.convbg_m = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
    self.bnbg_m = nn.BatchNorm2d(512)
    self.relubg_m = nn.ReLU(inplace=True)
    self.convbg_2 = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
    self.bnbg_2 = nn.BatchNorm2d(512)
    self.relubg_2 = nn.ReLU(inplace=True)

    self.conv6d_1 = nn.Conv2d(1024, 512, 3, padding=1)
    self.bn6d_1 = nn.BatchNorm2d(512)
    self.relu6d_1 = nn.ReLU(inplace=True)
    self.conv6d_m = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
    self.bn6d_m = nn.BatchNorm2d(512)
    self.relu6d_m = nn.ReLU(inplace=True)
    self.conv6d_2 = nn.Conv2d(512, 512, 3, dilation=2, padding=2)
    self.bn6d_2 = nn.BatchNorm2d(512)
    self.relu6d_2 = nn.ReLU(inplace=True)

    self.conv5d_1 = nn.Conv2d(1024, 512, 3, padding=1)
    self.bn5d_1 = nn.BatchNorm2d(512)
    self.relu5d_1 = nn.ReLU(inplace=True)
    self.conv5d_m = nn.Conv2d(512, 512, 3, padding=1)
    self.bn5d_m = nn.BatchNorm2d(512)
    self.relu5d_m = nn.ReLU(inplace=True)
    self.conv5d_2 = nn.Conv2d(512, 512, 3, padding=1)
    self.bn5d_2 = nn.BatchNorm2d(512)
    self.relu5d_2 = nn.ReLU(inplace=True)

    self.conv4d_1 = nn.Conv2d(1024, 512, 3, padding=1)
    self.bn4d_1 = nn.BatchNorm2d(512)
    self.relu4d_1 = nn.ReLU(inplace=True)
    self.conv4d_m = nn.Conv2d(512, 512, 3, padding=1)
    self.bn4d_m = nn.BatchNorm2d(512)
    self.relu4d_m = nn.ReLU(inplace=True)
    self.conv4d_2 = nn.Conv2d(512, 256, 3, padding=1)
    self.bn4d_2 = nn.BatchNorm2d(256)
    self.relu4d_2 = nn.ReLU(inplace=True)

    self.conv3d_1 = nn.Conv2d(512, 256, 3, padding=1)
    self.bn3d_1 = nn.BatchNorm2d(256)
    self.relu3d_1 = nn.ReLU(inplace=True)
    self.conv3d_m = nn.Conv2d(256, 256, 3, padding=1)
    self.bn3d_m = nn.BatchNorm2d(256)
    self.relu3d_m = nn.ReLU(inplace=True)
    self.conv3d_2 = nn.Conv2d(256, 128, 3, padding=1)
    self.bn3d_2 = nn.BatchNorm2d(128)
    self.relu3d_2 = nn.ReLU(inplace=True)

    self.conv2d_1 = nn.Conv2d(256, 128, 3, padding=1)
    self.bn2d_1 = nn.BatchNorm2d(128)
    self.relu2d_1 = nn.ReLU(inplace=True)
    self.conv2d_m = nn.Conv2d(128, 128, 3, padding=1)
    self.bn2d_m = nn.BatchNorm2d(128)
    self.relu2d_m = nn.ReLU(inplace=True)
    self.conv2d_2 = nn.Conv2d(128, 64, 3, padding=1)
    self.bn2d_2 = nn.BatchNorm2d(64)
    self.relu2d_2 = nn.ReLU(inplace=True)

    self.conv1d_1 = nn.Conv2d(128, 64, 3, padding=1)
    self.bn1d_1 = nn.BatchNorm2d(64)
    self.relu1d_1 = nn.ReLU(inplace=True)
    self.conv1d_m = nn.Conv2d(64, 64, 3, padding=1)
    self.bn1d_m = nn.BatchNorm2d(64)
    self.relu1d_m = nn.ReLU(inplace=True)
    self.conv1d_2 = nn.Conv2d(64, 64, 3, padding=1)
    self.bn1d_2 = nn.BatchNorm2d(64)
    self.relu1d_2 = nn.ReLU(inplace=True)

    self.upscore6 = nn.Upsample(
        scale_factor=32, mode="bilinear", align_corners=False
    )
    self.upscore5 = nn.Upsample(
        scale_factor=16, mode="bilinear", align_corners=False
    )
    self.upscore4 = nn.Upsample(
        scale_factor=8, mode="bilinear", align_corners=False
    )
    self.upscore3 = nn.Upsample(
        scale_factor=4, mode="bilinear", align_corners=False
    )
    self.upscore2 = nn.Upsample(
        scale_factor=2, mode="bilinear", align_corners=False
    )

    self.outconvb = nn.Conv2d(512, 1, 3, padding=1)
    self.outconv6 = nn.Conv2d(512, 1, 3, padding=1)
    self.outconv5 = nn.Conv2d(512, 1, 3, padding=1)
    self.outconv4 = nn.Conv2d(256, 1, 3, padding=1)
    self.outconv3 = nn.Conv2d(128, 1, 3, padding=1)
    self.outconv2 = nn.Conv2d(64, 1, 3, padding=1)
    self.outconv1 = nn.Conv2d(64, 1, 3, padding=1)

    self.refunet = _RefUnet(1, 64)

forward 

forward(
    pixel_values: Float[
        Tensor, "batch channels height width"
    ],
    return_dict: bool | None = None,
) -> (
    BASNetSaliencyOutput
    | tuple[Shaped[torch.Tensor, "..."], ...]
)

Predict saliency with the BASNet forward path.

Source code in models/basnet/src/basnet/modeling_basnet.py 
347
348
349
350
351
352
353
354
355
356
357
358
359
360
def forward(
    self,
    pixel_values: Float[torch.Tensor, "batch channels height width"],
    return_dict: bool | None = None,
) -> BASNetSaliencyOutput | tuple[Shaped[torch.Tensor, "..."], ...]:
    """Predict saliency with the BASNet forward path."""
    return_dict = (
        return_dict if return_dict is not None else self.config.use_return_dict
    )
    outputs = self._forward_impl(pixel_values)
    saliency = normalize_saliency(outputs[0][:, 0, :, :])
    if not return_dict:
        return (saliency, *outputs)
    return BASNetSaliencyOutput(saliency=saliency, side_outputs=outputs)

normalize_saliency 

normalize_saliency(
    pred: Float[Tensor, "... height width"],
) -> Float[torch.Tensor, "... height width"]

Normalize saliency maps independently over each spatial map.

Source code in models/basnet/src/basnet/modeling_basnet.py 
149
150
151
152
153
154
155
def normalize_saliency(
    pred: Float[torch.Tensor, "... height width"],
) -> Float[torch.Tensor, "... height width"]:
    """Normalize saliency maps independently over each spatial map."""
    min_value = pred.amin(dim=(-2, -1), keepdim=True)
    max_value = pred.amax(dim=(-2, -1), keepdim=True)
    return (pred - min_value) / (max_value - min_value)