How oassal/robio2024 works
ROBIO trains and evaluates a VAE student to estimate weld-pool shape from RGB images. The main workflow loads images and manual point annotations, derives teacher heatmaps from a segmentation model, trains the student by distillation across folds, then evaluates predicted points against teacher and manual annotations. The repository also supports DeepLab and YOLO segmentation paths, heatmap/point analysis, and optional evaluation video output. Some script and utility excerpts are partial, so the graph shows only wiring supported by the README or sampled source.
Main entry points: KD_Kfold_train.py K-fold training · eval.py Evaluation entry · models_teacher.py Teacher wrapper · dataset_kd.py KD dataset
Components 19
Data and targets
- KD dataset
dataset_kd.py
Teacher generation
- Teacher wrapper
models/models_teacher.py - DeepLab segmenter
models/SegModel.py - YOLO mask extraction
utils/yolo_functions.py - Heatmap generation
utils/utils_new.py - Contour geometry
utils/normals.py
Training and evaluation
- K-fold training
KD_Kfold_train.py - Evaluation entry
eval.py - Fold splits
trainers/kfold_functions.py - Training step
trainers/trainers.py - Distillation losses
losses.py
Model and analysis
- ResNet VAE
models/models_student.py - Heatmap metrics
trainers/heatmap_analysis.py
Other components
- Experimenter
- RGB images and labels
- Model checkpoints
- TensorBoard logs
- Evaluation video
- YOLO model
Connections 28
- Experimenter → K-fold training: starts training
- Experimenter → Evaluation entry: starts evaluation
- RGB images and labels → KD dataset: loads samples
- K-fold training → Fold splits: creates folds
- K-fold training → KD dataset: builds datasets
- K-fold training → DeepLab segmenter: loads segmenter
- KD dataset → DeepLab segmenter: runs segmentation
- KD dataset → Heatmap generation: derives targets
- DeepLab segmenter → Heatmap generation: provides masks
- Heatmap generation → Contour geometry: uses contour points
- K-fold training → Teacher wrapper: initializes teacher
- Teacher wrapper → DeepLab segmenter: runs segmentation
- Teacher wrapper → Heatmap generation: generates heatmaps
- Teacher wrapper → YOLO mask extraction: extracts masks
- YOLO mask extraction → YOLO model: runs tracking
- K-fold training → ResNet VAE: initializes student
- K-fold training → Training step: runs training
- Training step → ResNet VAE: updates weights
- Training step → Teacher wrapper: gets targets
- Training step → Distillation losses: computes loss
- K-fold training → TensorBoard logs: writes metrics
- Evaluation entry → KD dataset: loads evaluation data
- Evaluation entry → Model checkpoints: loads weights
- Evaluation entry → Teacher wrapper: uses teacher
- Evaluation entry → ResNet VAE: runs inference
- Evaluation entry → Heatmap metrics: measures distances
- Evaluation entry → TensorBoard logs: writes metrics
- Evaluation entry → Evaluation video: writes optional video