rsh-vibroarthrography-cnn

This project was created as part of the SEKMO program in collaboration with the Centre for Environmental Intelligence and Sensing at Taltech

This repository is an experiment workspace for classifying knee-joint vibroarthrography (VAG) signals with image-based deep learning workflows.

This repository is a work in progress and should be treated as experimental.

Current Status

Simple result summary from the current cross-validation pipeline:

• Grayscale mode

  • mean accuracy: 0.5614 +/- 0.0268
  • pathological F1: 0.3257 +/- 0.2231

• RGB mode

  • mean accuracy: 0.5941 +/- 0.0742
  • pathological F1: 0.4141 +/- 0.1473

Current conclusion:

• RGB performs better than grayscale under the same evaluation protocol.
• Both modes are still weak and unstable on this dataset.
• The project is not yet at a level where strong classification claims are justified.

Immediate next steps:

• Add class weighting or weighted sampling to improve pathological recall.
• Try reducing model capacity or freezing more ResNet layers.
• Add a patient-level or external holdout if leakage risk needs to be ruled out.

At a high level, the project:

1. Loads raw VAG signals from the OpenVAG dataset.
2. Preprocesses the signals.
3. Converts them into time-frequency images.
4. Trains CNN-based models on those derived images.

The current active training entrypoint is:

• scripts/training/train_resnet_cwt.py

Older experiment variants are kept in:

• legacy_entrypoints/

Repository Layout

• vag_cnn/

  • Reusable Python modules for loading data, preprocessing, CEEMDAN decomposition, plotting, time-frequency analysis, and PyTorch dataset/model helpers.

• scripts/preprocessing/

  • Scripts that generate image datasets from the raw VAG signals.

• scripts/training/

  • Model training scripts.

• legacy_entrypoints/

  • Older training and evaluation entrypoints preserved for reference.

• tests/

  • Small exploratory and plotting scripts. These are not a formal automated test suite.

• data/

  • Raw OpenVAG source data.

• images/

  • Generated image datasets produced by preprocessing scripts. These directories can be regenerated from the raw data in data/open_vag/.

• models/

  • Saved model weights.

• notebooks/

  • Exploratory notebook work.

Data Directories

data/

This contains the raw OpenVAG dataset under data/open_vag/.

• data/open_vag/normal/

  • 51 raw text files for healthy subjects.

• data/open_vag/pathology/

  • 38 raw text files for pathological subjects.

Total raw signals: 89.

The bundled data/open_vag/README.txt describes the source dataset as knee-joint sound recordings from 51 normal subjects and 38 pathological subjects.

Generated Image Directories

The image directories under images/ are not primary source data. They are generated artifacts created by the preprocessing scripts from the raw VAG signals in data/open_vag/.

Script-to-output mapping:

• scripts/preprocessing/create_cwt_images_ceemdan.py

  • reads raw healthy/pathological signals from data/open_vag/
  • preprocesses each signal
  • reconstructs a mid-band CEEMDAN signal
  • computes a continuous wavelet transform
  • writes 244x244 multichannel CWT PNGs to images/EMD_CWT_244/

• scripts/preprocessing/create_spectrograms.py

  • reads raw healthy/pathological signals from data/open_vag/
  • preprocesses each signal directly, without CEEMDAN reconstruction
  • computes STFT spectrogram images
  • writes PNGs to images/SPECS/

• scripts/preprocessing/create_spectrograms_ceemdan.py

  • reads raw healthy/pathological signals from data/open_vag/
  • preprocesses each signal
  • reconstructs a mid-band CEEMDAN signal
  • computes STFT spectrogram images from the reconstructed signal
  • writes PNGs to images/EMD_SPECS/

• scripts/preprocessing/create_tfds.py

  • reads raw healthy/pathological signals from data/open_vag/
  • preprocesses each signal
  • computes CEEMDAN IMFs and a reconstructed signal
  • computes time-frequency distribution images for each IMF and the reconstructed signal
  • writes PNGs to images/TFDs/

The repository currently contains several pre-generated image directories for convenience, but they can be rebuilt by re-running the corresponding preprocessing scripts.

images/EMD_CWT/

Legacy CEEMDAN-reconstructed continuous wavelet transform image set.

• Contains 89 PNG images total.
• Organized as healthy/ and pathology/.

This directory holds older CWT-derived images used by earlier training scripts.

images/EMD_CWT_244/

Current CWT image dataset used by the latest ResNet training script.

• Contains 89 PNG images total.
• Organized as healthy/ and pathology/.
• Generated by scripts/preprocessing/create_cwt_images_ceemdan.py.

This directory contains regenerated 244x244 multichannel CWT images built from CEEMDAN-reconstructed signals.

images/EMD_SPECS/

CEEMDAN-reconstructed spectrogram image dataset.

• Contains 89 PNG images total.
• Organized as healthy/ and pathology/.
• Generated by scripts/preprocessing/create_spectrograms_ceemdan.py.

These are spectrograms computed after CEEMDAN reconstruction.

images/SPECS/

Plain spectrogram image dataset created directly from preprocessed signals without CEEMDAN reconstruction.

• Organized as healthy/ and pathology/.
• Generated by scripts/preprocessing/create_spectrograms.py.

images/TFDs/

Time-frequency distribution image dataset for CEEMDAN IMFs and reconstructed signals.

• Organized as healthy/ and pathology/.
• Generated by scripts/preprocessing/create_tfds.py.
• Contains multiple images per original signal, including IMF-specific images and a reconstructed-signal image.

Current Workflow

The most current top-level path in the repository is:

1. Start from raw signals in data/open_vag/.
2. Generate CWT images with scripts/preprocessing/create_cwt_images_ceemdan.py.
3. Train a ResNet classifier with scripts/training/train_resnet_cwt.py.

That training script currently reads from:

• images/EMD_CWT_244/

and saves the best-fold weights to:

• models/resnet_spectrogram_classifier_cv_best_fold.pth
• models/resnet_spectrogram_classifier_cv_metrics.json
• models/resnet_spectrogram_classifier_cv_oof_predictions.csv

Evaluation Methodology

The canonical training script uses a fixed evaluation procedure intended to be more reliable than the previous single random holdout.

Current methodology in scripts/training/train_resnet_cwt.py:

1. A fixed random seed is used for Python, NumPy, and PyTorch.
2. The dataset is evaluated with stratified cross-validation.
3. Each fold preserves the healthy/pathological class balance as closely as possible.
4. The script supports both grayscale input and RGB input decoded from the same PNG files.
5. Training augmentation is limited to resizing plus random horizontal flip.
6. Validation folds use deterministic preprocessing only.
7. RGB mode uses ImageNet normalization; grayscale mode does not.
8. Metrics are computed fold by fold and then summarized with mean and standard deviation.
9. A final classification report is produced from aggregated out-of-fold predictions across all 89 samples.

Default values:

• seed=42
• input_mode=grayscale
• folds=5
• batch_size=16
• epochs=30
• lr=1e-4
• data_root=images/EMD_CWT_244/
• model_output=models/resnet_spectrogram_classifier_cv_best_fold.pth
• metrics_output=models/resnet_spectrogram_classifier_cv_metrics.json
• predictions_output=models/resnet_spectrogram_classifier_cv_oof_predictions.csv

This is still not the same as having a fully independent external test set, but it is materially better than reporting one favorable train/validation split from a small dataset.

How To Run

Run all commands from the repository root.

1. Create a virtual environment

python3 -m venv .venv
source .venv/bin/activate

The repository .gitignore already excludes .venv/.

2. Install dependencies

Install the project dependencies into the virtual environment:

pip install -r requirements.txt

3. Generate the current CWT dataset

python3 scripts/preprocessing/create_cwt_images_ceemdan.py

This reads raw files from data/open_vag/ and writes images into images/EMD_CWT_244/.

4. Train the current classifier

python3 scripts/training/train_resnet_cwt.py

This trains a ResNet18 classifier on images/EMD_CWT_244/, runs stratified 5-fold cross-validation, prints fold metrics plus an aggregated out-of-fold classification report, and writes the best-fold model weights to models/.

Artifacts written by the canonical training run:

• models/resnet_spectrogram_classifier_cv_best_fold.pth

  • best-fold model weights

• models/resnet_spectrogram_classifier_cv_metrics.json

  • run configuration, per-fold metrics, and aggregate metrics

• models/resnet_spectrogram_classifier_cv_oof_predictions.csv

  • one row per sample with image path, true label, and out-of-fold predicted label

Useful CLI options:

python3 scripts/training/train_resnet_cwt.py \
  --data-root images/EMD_CWT_244/ \
  --input-mode grayscale \
  --folds 5 \
  --epochs 30 \
  --batch-size 16 \
  --lr 1e-4 \
  --seed 42 \
  --model-output models/resnet_spectrogram_classifier_cv_best_fold.pth \
  --metrics-output models/resnet_spectrogram_classifier_cv_metrics.json \
  --predictions-output models/resnet_spectrogram_classifier_cv_oof_predictions.csv

To compare against the older RGB idea under the same evaluation protocol:

python3 scripts/training/train_resnet_cwt.py --input-mode rgb

In RGB mode, the script automatically writes separate artifacts with _rgb appended to the filename so they do not overwrite the grayscale outputs.

To inspect the full CLI:

python3 scripts/training/train_resnet_cwt.py --help

Full setup and run sequence:

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
python3 scripts/preprocessing/create_cwt_images_ceemdan.py
python3 scripts/training/train_resnet_cwt.py \
  --data-root images/EMD_CWT_244/ \
  --input-mode grayscale \
  --folds 5 \
  --epochs 30 \
  --batch-size 16 \
  --lr 1e-4 \
  --seed 42 \
  --model-output models/resnet_spectrogram_classifier_cv_best_fold.pth \
  --metrics-output models/resnet_spectrogram_classifier_cv_metrics.json \
  --predictions-output models/resnet_spectrogram_classifier_cv_oof_predictions.csv

Other Useful Scripts

• scripts/preprocessing/create_spectrograms.py

  • Builds plain spectrogram images in images/SPECS/.

• scripts/preprocessing/create_spectrograms_ceemdan.py

  • Builds CEEMDAN-based spectrogram images in images/EMD_SPECS/.

• scripts/preprocessing/create_tfds.py

  • Builds time-frequency distribution images in images/TFDs/.

• scripts/training/pytorch_train_pipeline.py

  • Autoencoder-based workflow on grouped IMF image sets.

• scripts/training/pytorch_train_pipeline_for_spectrograms.py

  • Autoencoder-based workflow for spectrogram images.

Important Caveat

This repository is still an experiment workspace, not a polished training package.

In particular:

• There are multiple legacy approaches preserved in the tree.
• The dataset is still small at 89 total signals.
• Cross-validation is better than a single split, but it is not a substitute for a true external test set.