# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch>
#
# SPDX-License-Identifier: GPL-3.0-or-later

"""CNN for Tuberculosis Detection.

Implementation of the model architecture proposed by F. Pasa in the article
"Efficient Deep Network Architectures for Fast Chest X-Ray Tuberculosis
Screening and Visualization".

Reference: [PASA-2019]_
"""

from torch import empty
from torch.nn import BCEWithLogitsLoss

from ...data.transforms import ElasticDeformation
from ...models.pasa import PASA

# config
optimizer_configs = {"lr": 8e-5}

# optimizer
optimizer = "Adam"

# criterion
criterion = BCEWithLogitsLoss(pos_weight=empty(1))
criterion_valid = BCEWithLogitsLoss(pos_weight=empty(1))

augmentation_transforms = [ElasticDeformation(p=0.8)]

# model
model = PASA(
    criterion,
    criterion_valid,
    optimizer,
    optimizer_configs,
    augmentation_transforms=augmentation_transforms,
)