A PyTorch pipeline that segments the wound region on a 3D facial scan using TSGCNet, a two-stream graph convolutional network — one stream reasons over triangle positions, the other over surface orientation, fused into a per-triangle wound/normal-skin prediction. Reproduces the method in Nguyen et al. (2024), re-engineered into a resumable, memory-efficient pipeline that trains on an ordinary laptop CPU — no GPU required.
Reproduction / engineering project, not novel research — see Background for attribution.
- Reads a 3D face mesh (thousands of connected triangles) with the wound region colour-coded in the training data.
- Computes a 24-number geometric descriptor per triangle (position, centre, and orientation vectors).
- Feeds it through TSGCNet to predict, per triangle, wound vs. normal skin.
- Extracts the wound region as its own printable 3D mesh (Algorithm 2: combines the segmented wound region with a separately-reconstructed "healed" face to determine the filling shape).
- Visualizes predictions and training metrics.
Figure adapted from Nguyen et al. (2024), Fig. 2 (CC BY 4.0). Implemented in wound_segmentation/model.py, verified layer-by-layer against this diagram — including the 64+128+256=448 concatenated feature width and the 512/1024-channel fusion stage shown above.
- Severe class imbalance (wounds cover a tiny fraction of each face) — compared four imbalance-aware loss functions (focal, Dice, weighted/plain cross-entropy) with a proper from-scratch-per-loss training protocol, so the comparison is actually meaningful.
- CPU-only training — redesigned the model's memory footprint (narrower channels, smaller graph neighbourhood, gradient accumulation) to fit a laptop instead of the reference GPU setup.
- Crash-resilient training loop — per-epoch checkpointing and CSV logging so a training run surviving a forced OS restart just resumes instead of losing hours of progress.
- Mesh geometry from scratch — sub-mesh extraction, mesh merging, and watertightness cleanup for the 3D-printable output, implemented and unit-tested in plain NumPy.
- Clean packaging — installable Python package + CLI scripts + config system, not just a notebook.
80 train / 20 test meshes, CPU-only laptop (i5-1240P), lite model config:
| Loss function | Val. accuracy | Val. mIoU |
|---|---|---|
| CrossEntropy | 0.9430 | 0.4715 |
| WeightedCrossEntropy | 0.9220 | 0.4788 |
| DiceLoss | 0.9430 | 0.4715 |
| FocalLoss | 0.9143 | 0.5429 |
(Early-run numbers — see experiment_fixed_nqd/laptop_run/results_loss_functions.csv
for the latest, and scripts/plot_training_history.py to chart them.)
For scale: the paper's own 100/20-mesh ablation (full GPU model, 50
epochs) reports 97.69% accuracy — a fair target here; its headline
0.9999993% figure used their full ~32,000-mesh dataset on a
workstation GPU.
pip install -r requirements.txt
# Train (CPU-friendly defaults; resumes automatically if interrupted)
python train.py --data-train-dir data/train --data-test-dir data/test --num-epochs 30
# Visualize a mesh's predicted wound region
python scripts/visualize_mesh.py --mesh sample_data/face_0.ply --checkpoint <path_to_checkpoint.pth> --out docs/comparison.png
# Chart training curves / loss-function comparison
python scripts/plot_training_history.py --history <path_to_results.csv> --out-dir docs/Full options: python train.py --help. Data layout: data/train/*.ply,
data/test/*.ply (wound faces coloured yellow, normal gray).
wound_segmentation/ # installable package: model, losses, dataset,
# metrics, filling extraction, visualization, config
train.py # CLI training entry point (resumable)
scripts/ # visualize_mesh.py, plot_training_history.py
sample_data/face_0.ply # small example mesh for demos
notebooks/ # exploratory notebooks
tests/ # pytest unit tests
- GitHub Actions CI (tests + smoke-test on push)
- More unit test coverage
- Wound-filling extraction (Algorithm 2)
- Segmentation + training-history visualizations
Method and architecture from: Nguyen et al., 2024, CMES (wound-filling pipeline) and Zhao et al., 2022, IEEE TMI (TSGCNet). Original reference implementation: SIMOGroup/WoundFilling3D. This repo is an independent, from-scratch reproduction built as an academic project — not affiliated with the original authors.
The code in this repository is MIT-licensed (see LICENSE) — this covers the code I wrote here, not the paper or method itself, which remain the original authors' published work (properly cited above). Reproducing a published method for learning/academic purposes is standard practice and doesn't require the authors' permission, provided it's clearly attributed and not presented as their work.