Skip to content

Latest commit

 

History

75 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

ComputerVision2026

Computer Vision course project for the AI & Robotics master's programme.

Project topic: anatomy-aware periodontal radiograph analysis on the DenPAR dataset.

This repository implements a complete pipeline for dental panoramic radiographs: dataset inspection, preprocessing, anatomical target construction, RBL-derived image-level severity labelling, U-Net baselines, anatomical autoencoder pretraining, and two anatomically constrained models inspired by ACNN-style priors.

The core idea is to use anatomical annotations not only as direct segmentation targets, but also as a learned prior over plausible dental anatomy. We first train a multi-task U-Net to predict anatomical maps and an image-level severity class. We then train an anatomical autoencoder on ground-truth anatomical maps only. The frozen autoencoder is used in two ways:

  1. ACNN-Soft: the autoencoder encoder defines a latent-space regularization loss.
  2. ACNN-Strong: the autoencoder encoder-decoder acts as a hard projection module for the predicted anatomical maps.

Important clinical note
The classification target used in this repository is an RBL-derived image-level severity class, computed from estimated radiographic bone loss (RBL). It is a project-specific image-level target and should not be interpreted as a full clinical periodontal staging label.


Team members

  • Livio Marzio della Penna
  • Claudia Julia Istoc

Dataset

This project uses DenPAR, a dental panoramic radiograph dataset with anatomical annotations and periodontal information.

The dataset is not included in this repository. Download it manually and extract it in the repository root, preserving the expected DenPAR folder structure used by the notebooks.


Repository structure

The repository is organized as a notebook-based experimental pipeline plus a reusable PyTorch codebase.

ComputerVision2026/
├── README.md
│
├── Dataset/                                # downloaded DenPAR dataset (not included)
│   ├── Testing/
│   ├── Training/
│   ├── Validation/
│   └── Characteristics of radiographs included.xlsx
│
├── notebooks/
│   ├── 01_dataset_exploration.ipynb
│   ├── 02_preprocessing.ipynb
│   ├── 02A_pretrain_sanity_check.ipynb
│   │
│   ├── 03_train_unet_baseline.ipynb
│   ├── 03A_train_unet_baseline_v2.ipynb
│   ├── 03B_train_unet_baseline_v3.ipynb
│   │
│   ├── 04_train_autoencoder.ipynb
│   ├── 05_soft_constraint.ipynb
│   ├── 07_hard_constraint.ipynb
│   │
│   └── 06_simple_hard_constraint.ipynb     # archived / latent-fusion prototype
│
├── src/
│   ├── data/
│   │   └── data.py
│   ├── models/
│   │   ├── blocks.py
│   │   ├── unet.py
│   │   ├── autoencoder.py
│   │   └── hard.py
│   └── training/
│       ├── losses.py
│       └── metrics.py
│
├── Processed_Data/                         # generated by 02_preprocessing.ipynb
│   └── processed_512x512/
│       ├── processed_labels_final_clean.csv
│       ├── preprocessing_config.json
│       ├── rbl_detail_records.csv
│       ├── rbl_label_diagnostics.csv
│       ├── images/
│       └── anatomical_maps/
│
├── runs/                                   # generated by training notebooks
│   ├── unet_baseline_v1/
│   ├── unet_baseline_v2/
│   ├── unet_baseline_v3/
│   ├── anatomical_autoencoder_v2/
│   ├── acnn_soft_v1/
│   └── acnn_hard_v1/
│
└── References/                             # Material references and additional readings and project notes

The notebooks assume that the src package is importable from the project root.


Environment setup

The project was developed with Python 3.11 and PyTorch. A CUDA-enabled GPU is strongly recommended for the training notebooks.

conda create -n cv26 python=3.11 -y
conda activate cv26

python -m pip install --upgrade pip
pip install ipykernel jupyterlab
python -m ipykernel install --user --name cv26 --display-name "Python (CV 2026)"

# CUDA 12.1 PyTorch build. Change this command if your CUDA/PyTorch setup differs.
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

pip install numpy pandas matplotlib pillow opencv-contrib-python scikit-image scikit-learn tqdm openpyxl pyyaml pycocotools shapely rich

High-level pipeline

The recommended execution order is:

1. 01_dataset_exploration.ipynb
2. 02_preprocessing.ipynb
3. 02A_pretrain_sanity_check.ipynb
4. 03_train_unet_baseline.ipynb          # simple baseline, BCE segmentation
5. 03A_train_unet_baseline_v2.ipynb      # main U-Net baseline, channel-aware loss
6. 03B_train_unet_baseline_v3.ipynb      # sparse-channel ablation
7. 04_train_autoencoder.ipynb            # anatomical autoencoder prior
8. 05_soft_constraint.ipynb              # ACNN-Soft
9. 07_hard_constraint.ipynb              # ACNN-Strong

The main experimental comparison is:

U-Net baseline with channel-aware loss
vs
ACNN-Soft from scratch
vs
ACNN-Strong from scratch

The older notebook 06_simple_hard_constraint.ipynb is kept as development history / ablation prototype. It is not the primary hard-constraint implementation.


1. Dataset exploration

Notebook

notebooks/01_dataset_exploration.ipynb

Purpose

This notebook inspects the raw DenPAR data before any model training. It is used to understand:

  • the original image format;
  • the available anatomical annotations;
  • the train/validation/test split fields;
  • possible missing annotations;
  • the structure of keypoint, polygon, and line annotations;
  • the feasibility of deriving an image-level RBL class.

The output of this notebook is mostly diagnostic. It does not define the final model inputs; that happens in 02_preprocessing.ipynb.


2. Preprocessing and target construction

Notebook

notebooks/02_preprocessing.ipynb

Main output

The preprocessing notebook builds a cleaned 512×512 working dataset:

Processed_Data/processed_512x512/processed_labels_final_clean.csv
Processed_Data/processed_512x512/images/
Processed_Data/processed_512x512/anatomical_maps/

The generated CSV is the central metadata file used by all downstream notebooks.

Image preprocessing

Each radiograph is converted to grayscale, resized to 512×512, normalized to [0, 1], and saved as a processed image. The corresponding anatomical annotations are transformed consistently to the resized coordinate system.

Main constants:

TARGET_HEIGHT = 512
TARGET_WIDTH = 512
CEJ_SIGMA = 4.0
APEX_SIGMA = 4.0
BONE_LINE_THICKNESS = 3

Anatomical target maps

For each image, preprocessing creates an anatomical map tensor:

anatomical_map: [4, 512, 512]

The channel convention is fixed across the whole project:

Channel Name Target type Meaning
0 tooth_mask binary mask tooth region mask
1 cej_heatmap continuous heatmap cemento-enamel junction points
2 apex_heatmap continuous heatmap root apex points
3 bone_line_mask binary/thin mask alveolar bone-line annotation

CEJ and apex heatmaps

CEJ and apex annotations are point-like. Instead of predicting single pixels, the project converts them into Gaussian heatmaps. For a set of annotated points $P$, the heatmap value at pixel $(x, y)$ is:

$$H(x,y) = \max_{p_i \in P} \exp\left(-\frac{(x-x_i)^2 + (y-y_i)^2}{2\sigma^2}\right)$$

with $\sigma = 4$ for both CEJ and apex. The heatmap is clipped to [0, 1].

This makes the target smoother and more learnable than a one-pixel target, while still preserving the sparse anatomical localization task.

Bone-line mask

The bone-line annotations are rasterized as thin line masks with thickness 3. This gives a binary target for the alveolar bone line while preserving its thin structure.


RBL-derived image-level severity class

The project derives a classification label from estimated radiographic bone loss (RBL). The label is image-level, not tooth-level.

For each usable tooth, a CEJ point and an apex point define a local anatomical axis:

$$v = a - c$$

where:

  • $c$ is the CEJ point;
  • $a$ is the apex point;
  • $v$ is the CEJ-to-apex axis.

For each bone-line point $b$, the point is projected onto the CEJ-apex axis:

$$t = \frac{(b-c) \cdot v}{\lVert v \rVert^2}$$

The scalar $t$ is clipped to [0, 1]:

$$t_{clipped} = \min(1, \max(0, t))$$

In this convention:

  • $t=0$ means the projected bone point lies at the CEJ;
  • $t=1$ means the projected bone point lies at the apex.

The RBL estimate for a tooth is the selected projected value along the axis:

$$RBL_{tooth} = t_{clipped}$$

The image-level RBL value is the maximum tooth-level RBL estimate in the image:

$$RBL_{image} = \max_{tooth} RBL_{tooth}$$

The image-level class is then defined as:

Class Rule Interpretation used in this project
0 max_rbl < 0.15 low RBL
1 0.15 <= max_rbl <= 0.33 moderate RBL
2 max_rbl > 0.33 high RBL

The thresholds correspond to 15% and 33% of the CEJ-apex axis.

Again, this is an RBL-derived image-level severity class. It is not a complete clinical diagnosis.


Label reliability and filtering

The preprocessing notebook stores diagnostic flags for downstream filtering. Important fields include:

Column Meaning
max_rbl maximum estimated image-level RBL
rbl_class derived class 0/1/2
has_valid_label whether a class could be computed
has_reliable_label whether the label passed uncertainty checks
flag_invalid_label invalid or missing RBL-derived label
flag_saturated_rbl extremely high/saturated RBL estimate
flag_large_max_gap diagnostic flag based on RBL distribution within the image
flag_uncertain_rbl_label final uncertainty flag for classification
exclude_for_leakage whether the image is excluded due to duplicate/leakage detection
use_for_segmentation usable for anatomical segmentation
use_for_classification usable for image-level classification
use_for_multitask_training usable for joint segmentation/classification training

Duplicate and leakage control

To reduce train/validation/test leakage, the preprocessing includes a SIFT-based near-duplicate check. The strict duplicate rule uses:

SIFT_MIN_GOOD_MATCHES = 150
SIFT_MIN_INLIERS = 140
SIFT_MIN_INLIER_RATIO = 0.90
SIFT_MIN_COVERAGE_RATIO = 0.08

A pair of images is automatically considered strong duplicate/leakage evidence only when all strict criteria are satisfied. Confirmed near-duplicates are excluded through the exclude_for_leakage flag.


3. Dataset and dataloaders

Module

src/data/data.py

Main classes and functions

DenPARDataset
make_datasets
make_dataloaders
compute_class_weights
summarize_datasets

DenPARDataset reads processed_labels_final_clean.csv, loads the processed image and anatomical map, and returns a dictionary:

{
    "image": Tensor[1, H, W],
    "anat_map": Tensor[4, H, W],
    "label": Tensor[],
    "image_id": str,
    "split": str,
    "index": int,
    ... optional metadata ...
}

Supported task modes:

Mode Used by Filtering logic
autoencoder 04_train_autoencoder.ipynb anatomical maps only; segmentation-usable samples
segmentation optional segmentation-only experiments segmentation-usable samples
classification optional classification-only experiments valid/reliable classification labels
multitask U-Net, ACNN-Soft, ACNN-Strong samples usable for both segmentation and classification
all diagnostics minimal filtering

The main training notebooks use task_mode="multitask", while the anatomical autoencoder uses task_mode="autoencoder".


4. Shared model components

Module

src/models/blocks.py

The project uses a small set of reusable convolutional blocks:

Component Purpose
ConvBlock two convolutional layers with normalization and activation
UpBlock U-Net upsampling block with skip connection
UpBlockAE autoencoder upsampling block without skip connection
count_trainable_parameters utility for reporting model size

The U-Net decoder uses skip connections; the anatomical autoencoder decoder does not.


5. U-Net multi-task baseline

Module

src/models/unet.py

Notebooks

notebooks/03_train_unet_baseline.ipynb
notebooks/03A_train_unet_baseline_v2.ipynb
notebooks/03B_train_unet_baseline_v3.ipynb

Architecture

The main baseline is a multi-task U-Net. Given an input radiograph:

$$I \in \mathbb{R}^{1 \times 512 \times 512}$$

it predicts:

$$S_\theta(I) \in \mathbb{R}^{4 \times 512 \times 512}$$

for anatomical segmentation logits, and:

$$C_\theta(I) \in \mathbb{R}^{3}$$

for image-level class logits.

The model returns:

outputs = model(images)
seg_logits = outputs["seg_logits"]  # [B, 4, 512, 512]
cls_logits = outputs["cls_logits"]  # [B, 3]

Important implementation detail: seg_logits are raw logits, not probabilities. The loss uses BCEWithLogitsLoss; metrics and visualizations apply torch.sigmoid explicitly.

Multi-task objective

The standard multi-task objective is:

$$\mathcal{L}_{total} = \lambda_{seg}\mathcal{L}_{seg} + \lambda_{cls}\mathcal{L}_{cls}$$

where:

$$\lambda_{seg}=1, \qquad \lambda_{cls}=1$$

in the main experiments.

The classification loss is standard cross-entropy:

$$\mathcal{L}_{cls} = CE(C_\theta(I), y)$$

where $y \in {0,1,2}$ is the RBL-derived image-level class.


Channel-aware segmentation loss

The final U-Net baseline uses a channel-aware segmentation loss implemented in:

src/training/losses.py

The four anatomical channels have very different target densities. The tooth mask is dense, while CEJ, apex, and bone-line targets are sparse. A single uniform BCE loss is therefore not ideal.

The project uses:

Channel Loss terms
tooth mask weighted BCE + Dice
CEJ heatmap weighted BCE + MSE
apex heatmap weighted BCE + MSE
bone-line mask weighted BCE + Dice

For a binary/continuous target $M_c$ and logits $S_c$, the soft Dice term is:

$$Dice_c = \frac{2\sum \sigma(S_c)M_c + \epsilon}{\sum \sigma(S_c) + \sum M_c + \epsilon}$$ $$\mathcal{L}_{Dice,c} = 1 - Dice_c$$

For heatmap channels, the MSE term is computed on probabilities:

$$\mathcal{L}_{MSE,c} = \lVert \sigma(S_c) - M_c \rVert_2^2$$

The per-channel losses are:

$$\mathcal{L}_{tooth} = BCE_{pos}(S_{tooth}, M_{tooth}) + \alpha_{tooth}\mathcal{L}_{Dice,tooth}$$ $$\mathcal{L}_{cej} = BCE_{pos}(S_{cej}, M_{cej}) + \beta\mathcal{L}_{MSE,cej}$$ $$\mathcal{L}_{apex} = BCE_{pos}(S_{apex}, M_{apex}) + \beta\mathcal{L}_{MSE,apex}$$ $$\mathcal{L}_{bone} = BCE_{pos}(S_{bone}, M_{bone}) + \alpha_{bone}\mathcal{L}_{Dice,bone}$$

The final segmentation loss is a normalized weighted average:

$$\mathcal{L}_{seg} = \frac{ w_{tooth}\mathcal{L}_{tooth} + w_{cej}\mathcal{L}_{cej} + w_{apex}\mathcal{L}_{apex} + w_{bone}\mathcal{L}_{bone} }{w_{tooth}+w_{cej}+w_{apex}+w_{bone}}$$

All segmentation loss computations explicitly cast logits and targets to float32. This avoids FP16 overflow when AMP is enabled and losses/metrics sum over dense 512×512 maps.


U-Net baseline variants

03_train_unet_baseline.ipynb — v1 baseline

This is the first simple baseline. It trains the same multi-task U-Net architecture with simpler segmentation supervision. It is kept as a reference and development baseline.

03A_train_unet_baseline_v2.ipynb — main U-Net baseline

This is the main U-Net baseline used for comparison against ACNN variants. It uses the channel-aware loss with a balanced configuration:

seg_channel_weights = {
    "tooth_mask": 1.0,
    "cej_heatmap": 2.0,
    "apex_heatmap": 2.0,
    "bone_line_mask": 2.0,
}

seg_pos_weights = {
    "tooth_mask": 1.0,
    "cej_heatmap": 10.0,
    "apex_heatmap": 10.0,
    "bone_line_mask": 8.0,
}

tooth_dice_weight = 1.0
bone_dice_weight = 0.5
heatmap_mse_weight = 2.0

This configuration is used again in ACNN-Soft and ACNN-Strong to make the comparison fair: the anatomical constraint is the main changed variable, not the base segmentation loss.

03B_train_unet_baseline_v3.ipynb — sparse-channel ablation

This notebook uses the same U-Net architecture and the same general channel-aware loss structure, but changes the weights to test a sharper sparse-channel trade-off:

seg_channel_weights = {
    "tooth_mask": 1.0,
    "cej_heatmap": 1.5,
    "apex_heatmap": 2.5,
    "bone_line_mask": 1.5,
}

seg_pos_weights = {
    "tooth_mask": 1.0,
    "cej_heatmap": 6.0,
    "apex_heatmap": 12.0,
    "bone_line_mask": 5.0,
}

tooth_dice_weight = 1.0
bone_dice_weight = 0.25
heatmap_mse_weight = 4.0

This is not a new architecture. It is a loss-weight ablation and has richer checkpointing/model-selection logic.


6. Anatomical autoencoder prior

Module

src/models/autoencoder.py

Notebook

notebooks/04_train_autoencoder.ipynb

Purpose

The anatomical autoencoder learns a compact latent representation of plausible anatomical maps. It is trained only on ground-truth anatomical maps, not on radiographs.

This makes it an anatomical prior that can later be frozen and reused in ACNN-Soft and ACNN-Strong.

Architecture

The autoencoder maps anatomical maps to a latent vector and reconstructs them:

$$M \in \mathbb{R}^{4 \times 512 \times 512}$$ $$z = E_\phi(M), \qquad z \in \mathbb{R}^{d}$$ $$\hat{M}_{logits} = D_\psi(z), \qquad \hat{M}_{logits} \in \mathbb{R}^{4 \times 512 \times 512}$$

The implementation returns logits:

outputs = ae(anat_maps)
z = outputs["latent"]
reconstruction_logits = outputs["reconstruction"]

The final autoencoder configuration uses:

in_channels = 4
out_channels = 4
base_channels = 32
latent_dim = 256
up_mode = "bilinear"

Autoencoder loss

The autoencoder is trained with the same channel-aware segmentation loss used for the U-Net v2 baseline:

$$\mathcal{L}_{AE} = \mathcal{L}_{seg}(D_\psi(E_\phi(M)), M)$$

This makes the reconstruction objective aware of dense and sparse anatomical channels.

Checkpoints

The notebook saves multiple checkpoints, including:

best_val_loss.pt
best_val_sparse_soft_dice.pt
last_autoencoder.pt

The main prior used by ACNN models is selected from the best validation loss checkpoint. The sparse checkpoint is useful for diagnostics.


7. ACNN-Soft: latent-space anatomical regularization

Notebook

notebooks/05_soft_constraint.ipynb

Concept

ACNN-Soft keeps the U-Net segmentation output as the final output, but adds a soft anatomical regularization term in the latent space of the frozen anatomical autoencoder.

The frozen autoencoder encoder maps both predicted anatomical maps and ground-truth anatomical maps to latent codes. The model is penalized if these latent codes are far apart.

Forward path

Given an image $I$:

$$S_\theta(I) = \text{raw U-Net segmentation logits}$$ $$P_\theta(I) = \sigma(S_\theta(I))$$

where $P_\theta(I)$ is the predicted anatomical probability map.

The frozen anatomical encoder computes:

$$z_{pred} = E_\phi(P_\theta(I))$$

For the ground-truth map $M$:

$$z_{gt} = E_\phi(M)$$

The autoencoder parameters are frozen. Gradients flow through $z_{pred}$ back to the U-Net, but not through $z_{gt}$.

Latent loss

$$\mathcal{L}_{latent} = \lVert z_{pred} - z_{gt} \rVert_2^2$$

The full objective is:

$$\mathcal{L}_{total} = \lambda_{seg}\mathcal{L}_{seg}(S_\theta(I), M) + \lambda_{cls}\mathcal{L}_{cls}(C_\theta(I), y) + \lambda_{latent}^{(e)}\mathcal{L}_{latent}$$

where the latent weight is warmed up over the first epochs:

$$\lambda_{latent}^{(e)} = \lambda_{latent} \cdot \min\left(1, \frac{e}{E_{warmup}}\right)$$

Main configuration:

lambda_latent = 1e-4
latent_warmup_epochs = 10
latent_batch_size = 4

latent_batch_size computes the latent loss only on a subset of each batch. This reduces memory and time while preserving the anatomical regularization signal.

Important distinction from ACNN-Strong

In ACNN-Soft, the final segmentation prediction is still the U-Net output:

final segmentation = raw U-Net segmentation logits

The autoencoder is used only as an extra training loss. It does not modify the output map at inference time.


8. ACNN-Strong: hard anatomical projection

Module

src/models/hard.py

Notebook

notebooks/07_hard_constraint.ipynb

Concept

ACNN-Strong uses the frozen anatomical autoencoder as an explicit projection module. The U-Net first predicts raw anatomical logits. These logits are converted to probabilities and passed through the frozen anatomical autoencoder encoder-decoder. The projected output becomes the final segmentation prediction.

In short:

ACNN-Soft:
    penalize predictions that are far from the anatomical latent manifold

ACNN-Strong:
    force the final segmentation output through the anatomical autoencoder projection

Forward path

The model computes:

$$S^{raw}_\theta(I) = \text{U-Net raw segmentation logits}$$ $$P^{raw}_\theta(I) = \sigma(S^{raw}_\theta(I))$$

Then the frozen anatomical autoencoder projects the raw probability map:

$$z = E_\phi(P^{raw}_\theta(I))$$ $$S^{proj}_\theta(I) = D_\psi(z)$$

The final segmentation output is:

$$S^{final}_\theta(I) = S^{proj}_\theta(I)$$

The model returns both projected and raw logits:

outputs = model(images)
projected_seg_logits = outputs["seg_logits"]
raw_seg_logits = outputs["raw_seg_logits"]
cls_logits = outputs["cls_logits"]
latent = outputs["latent"]

Classification head

The strong model also uses the anatomical latent code for classification. The classification head receives:

  1. the U-Net bottleneck feature map;
  2. the anatomical latent vector extracted from the raw predicted maps.

After global average pooling of the bottleneck, the two vectors are normalized, concatenated, and classified:

$$h = GAP(B_\theta(I))$$ $$\hat{y} = Head([h, z])$$

This makes the classifier anatomy-aware, while the segmentation output is anatomically projected.

Strong loss

The main segmentation loss is computed on the projected logits:

$$\mathcal{L}_{proj} = \mathcal{L}_{seg}(S^{proj}_\theta(I), M)$$

A smaller auxiliary loss is also applied to the raw U-Net logits:

$$\mathcal{L}_{raw} = \mathcal{L}_{seg}(S^{raw}_\theta(I), M)$$

The auxiliary raw loss prevents the U-Net from learning arbitrary raw outputs that are only corrected by the autoencoder.

The final strong segmentation loss is:

$$\mathcal{L}_{seg}^{strong} = \mathcal{L}_{proj} + \gamma \mathcal{L}_{raw}$$

with:

raw_aux_seg_weight = 0.25

The full ACNN-Strong objective is:

$$\mathcal{L}_{total} = \lambda_{seg}\mathcal{L}_{seg}^{strong} + \lambda_{cls}\mathcal{L}_{cls}$$

Numerical stability

The autoencoder projection is executed in float32 even when AMP is enabled:

raw_seg_probs = torch.sigmoid(raw_seg_logits.float())

with torch.amp.autocast(device_type=raw_seg_probs.device.type, enabled=False):
    latent = anatomical_autoencoder.encode(raw_seg_probs)
    projected_seg_logits = anatomical_autoencoder.decode(latent)

This is intentional. During early training, random U-Net predictions can be far outside the distribution of ground-truth anatomical maps. Running the frozen autoencoder projection in FP16 can produce numerical instability or NaNs. Keeping only this projection in FP32 preserves stability while allowing the rest of the model to benefit from AMP.


9. Metrics and model selection

Module

src/training/metrics.py

Classification metrics

The notebooks report:

  • accuracy;
  • balanced accuracy;
  • macro precision;
  • macro recall;
  • macro F1;
  • weighted F1.

Macro F1 is especially important because the RBL-derived classes may be imbalanced.

Soft segmentation metrics

Soft Dice is computed on sigmoid probabilities without thresholding:

$$Dice_c^{soft} = \frac{2\sum \sigma(S_c)M_c + \epsilon}{\sum \sigma(S_c) + \sum M_c + \epsilon}$$

The mean soft Dice is:

$$MeanSoftDice = \frac{1}{4}\sum_{c=1}^{4} Dice_c^{soft}$$

The sparse soft Dice excludes the dense tooth channel:

$$SparseSoftDice = \frac{1}{3}(Dice_{cej}^{soft} + Dice_{apex}^{soft} + Dice_{bone}^{soft})$$

This metric is useful because the tooth mask is much easier and denser than CEJ, apex, and bone-line targets.

Hard segmentation metrics

For thresholded metrics, logits are converted to probabilities and then binarized:

$$\hat{M}_c = \mathbb{1}[\sigma(S_c) \geq \tau_c]$$

Evaluation thresholds:

seg_eval_thresholds = {
    "tooth_mask": 0.5,
    "cej_heatmap": 0.3,
    "apex_heatmap": 0.3,
    "bone_line_mask": 0.5,
}

The notebooks report per-channel:

  • hard Dice;
  • IoU;
  • precision;
  • recall;
  • F1;
  • predicted positive fraction;
  • target positive fraction.

The predicted/target positive fractions are useful diagnostics for sparse channels. They show whether the model is predicting everything positive, everything negative, or a plausible amount of foreground.

Combined checkpoint score

For the main model-selection checkpoint, the notebooks use a balanced score:

$$CombinedScore = 0.5 \cdot MacroF1 + 0.5 \cdot MeanSoftDice$$

This score is used only for checkpoint selection. It is not a training objective.

The usual checkpoint set is:

last_*.pt
best_val_loss.pt
best_val_cls_macro_f1.pt
best_val_mean_soft_dice.pt
best_val_sparse_soft_dice.pt
best_val_combined_score.pt

The main checkpoint for final test evaluation is usually:

best_val_combined_score.pt

because the task is multi-objective: both image-level classification and anatomical segmentation matter.


10. Training outputs

Each training notebook creates a run folder under runs/:

runs/<run_name>/
├── config.json
├── history.csv
├── run_summary.json
├── test_metrics.csv
├── test_classification_predictions.csv
├── test_confusion_matrix.csv
├── checkpoints/
└── plots/

Typical files:

File Meaning
config.json exact configuration used for the run
history.csv epoch-level training/validation metrics
run_summary.json selected checkpoint and final test summary
test_metrics.csv final test metrics
test_classification_predictions.csv image-level predictions and probabilities
test_confusion_matrix.csv classification confusion matrix
checkpoints/*.pt saved PyTorch checkpoints
plots/ diagnostic and qualitative plots

11. Notebook-by-notebook guide

01_dataset_exploration.ipynb

Explores the raw DenPAR dataset and annotations. Use this to verify that paths, annotation formats, and splits are understood before preprocessing.

02_preprocessing.ipynb

Builds the final processed dataset:

  • resizes radiographs to 512×512;
  • creates four-channel anatomical maps;
  • estimates RBL values through CEJ-apex axis projection;
  • derives image-level RBL classes;
  • computes reliability flags;
  • runs SIFT leakage checks;
  • writes the final cleaned CSV.

02A_pretrain_sanity_check.ipynb

Checks that the processed dataset can be loaded correctly:

  • image tensors have shape [1, 512, 512];
  • anatomical maps have shape [4, 512, 512];
  • class labels are available where expected;
  • train/validation/test splits are usable;
  • visualization of processed images and targets is sensible.

03_train_unet_baseline.ipynb

Initial U-Net baseline. Useful as a simple reference run and development notebook.

03A_train_unet_baseline_v2.ipynb

Main U-Net baseline. Uses the balanced channel-aware loss configuration. This is the primary conventional baseline for comparison against ACNN-Soft and ACNN-Strong.

03B_train_unet_baseline_v3.ipynb

Sparse-channel ablation. Same U-Net architecture, different sparse-channel loss weights and richer checkpointing logic.

04_train_autoencoder.ipynb

Trains the anatomical autoencoder on ground-truth anatomical maps. The resulting frozen autoencoder defines the anatomical prior used by ACNN-Soft and ACNN-Strong.

05_soft_constraint.ipynb

Trains ACNN-Soft from scratch:

  • U-Net predicts segmentation and class logits;
  • frozen AE encoder extracts predicted and ground-truth latent codes;
  • latent MSE is added as a soft anatomical regularizer;
  • final segmentation remains the raw U-Net output.

07_hard_constraint.ipynb

Trains ACNN-Strong from scratch:

  • U-Net predicts raw segmentation logits;
  • raw predictions are projected through the frozen AE encoder-decoder;
  • projected logits are used as final segmentation output;
  • raw logits receive auxiliary segmentation supervision;
  • classification uses both U-Net bottleneck features and anatomical latent features.

Archived/development notebook: 06_simple_hard_constraint.ipynb

Prototype / ablation for latent fusion. This is not the final hard projection model.


12. Practical notes and implementation decisions

Logits vs probabilities

All model segmentation outputs are logits. The code applies sigmoid only when needed:

  • inside BCE-compatible losses when required;
  • inside Dice/MSE terms;
  • inside metrics;
  • before feeding predicted maps to the anatomical autoencoder.

This avoids accidental double-sigmoid or mismatched loss inputs.

Why CEJ and apex use heatmaps

CEJ and apex are sparse point-like annotations. Training with one-pixel binary maps would make the positive class extremely rare and unstable. Gaussian heatmaps provide a smoother target and make localization learnable.

Why sparse-channel diagnostics matter

The tooth channel can dominate mean metrics because it is dense and easier. For this reason, the notebooks also track sparse-channel performance separately:

CEJ + apex + bone-line

This prevents the model from looking strong only because it segments teeth well.

Why the anatomical autoencoder is frozen

The autoencoder is trained before the ACNN models. During ACNN training it is frozen because it represents a fixed anatomical prior. If it were updated jointly, the prior itself could move to accommodate bad predictions, weakening the constraint.

Why ACNN-Strong keeps a raw auxiliary loss

Without raw auxiliary supervision, the U-Net could learn raw maps that are not meaningful by themselves, relying entirely on the autoencoder decoder to produce reasonable outputs. The raw auxiliary term keeps the intermediate prediction grounded:

$$\mathcal{L}_{seg}^{strong} = \mathcal{L}_{proj} + 0.25\mathcal{L}_{raw}$$

Why parts of the code force float32

AMP is useful for speed and memory, but dense Dice-style reductions over large 512×512 maps can overflow or become unstable in FP16. The loss and metric implementations cast logits and targets to float32. ACNN-Strong also runs the frozen AE projection in float32 to avoid NaNs during early training.


13. Reproducibility checklist

To reproduce the full project from scratch:

  1. Create the environment.
  2. Download and extract DenPAR into the repository root.
  3. Run 01_dataset_exploration.ipynb for inspection.
  4. Run 02_preprocessing.ipynb to generate the processed dataset.
  5. Run 02A_pretrain_sanity_check.ipynb to verify processed data.
  6. Train the U-Net baseline with 03A_train_unet_baseline_v2.ipynb.
  7. Optionally train the sparse-channel ablation with 03B_train_unet_baseline_v3.ipynb.
  8. Train the anatomical autoencoder with 04_train_autoencoder.ipynb.
  9. Train ACNN-Soft with 05_soft_constraint.ipynb.
  10. Train ACNN-Strong with 07_hard_constraint.ipynb.
  11. Use the saved run_summary.json, test_metrics.csv, prediction CSVs, and qualitative plots for the report.

14. Final Summary

Model Main idea Segmentation output Anatomical prior usage
U-Net baseline standard multi-task model raw U-Net logits none
ACNN-Soft latent-space regularization raw U-Net logits frozen AE encoder in loss
ACNN-Strong hard anatomical projection AE-projected logits frozen AE encoder-decoder in forward path

Formally, the three main models can be summarized as:

U-Net baseline:
    learn anatomy directly from radiograph supervision

ACNN-Soft:
    learn anatomy from radiographs while being softly pulled toward the anatomical latent manifold

ACNN-Strong:
    force the final anatomical output to pass through the learned anatomical manifold

15. Naming conventions used in this repository

To avoid ambiguity:

  • RBL-derived severity class means the project-specific image-level class computed from estimated RBL.
  • Anatomical maps means the four-channel target tensor: tooth, CEJ, apex, bone-line.
  • Sparse channels means CEJ, apex, and bone-line.
  • ACNN-Soft means latent-space regularization through the frozen AE encoder.
  • ACNN-Strong means hard projection through the frozen AE encoder-decoder.
  • Projected segmentation means the ACNN-Strong output after autoencoder projection.
  • Raw segmentation means the intermediate U-Net output before autoencoder projection.

About

Group Project for CV course (AIRO)

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages