# oassal/robio2024 architecture

> Architecture diagram of the GitHub repository oassal/robio2024, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-27). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/oassal/robio2024
- Repository: https://github.com/oassal/robio2024

## Overview

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`](https://github.com/oassal/robio2024/blob/main/KD_Kfold_train.py) (K-fold training), [`eval.py`](https://github.com/oassal/robio2024/blob/main/eval.py) (Evaluation entry), [`models/models_teacher.py`](https://github.com/oassal/robio2024/blob/main/models/models_teacher.py) (Teacher wrapper), [`dataset_kd.py`](https://github.com/oassal/robio2024/blob/main/dataset_kd.py) (KD dataset)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Data and targets"]
  node_dataset["KD dataset<br/>[dataset_kd.py]"]
end

subgraph group_teacher["Teacher generation"]
  node_teacher["Teacher wrapper<br/>[models_teacher.py]"]
  node_segmenter["DeepLab segmenter<br/>[SegModel.py]"]
  node_yolomasks["YOLO mask extraction<br/>[yolo_functions.py]"]
  node_heatmaps["Heatmap generation<br/>[utils_new.py]"]
  node_points["Contour geometry<br/>[normals.py]"]
end

subgraph group_learning["Training and evaluation"]
  node_foldtrain["K-fold training<br/>[KD_Kfold_train.py]"]
  node_evalentry["Evaluation entry<br/>[eval.py]"]
  node_splits["Fold splits<br/>[kfold_functions.py]"]
  node_trainer["Training step<br/>[trainers.py]"]
  node_losses["Distillation losses<br/>[losses.py]"]
end

subgraph group_model["Model and analysis"]
  node_student["ResNet VAE<br/>[models_student.py]"]
  node_analysis["Heatmap metrics"]
end

node_user(("Experimenter"))
node_rgbdata["RGB images and labels"]
node_checkpoints["Model checkpoints"]
node_logs["TensorBoard logs"]
node_video["Evaluation video"]
node_yolo["YOLO model"]

node_user -->|"starts training"| node_foldtrain
node_user -->|"starts evaluation"| node_evalentry
node_rgbdata -->|"loads samples"| node_dataset
node_foldtrain -->|"creates folds"| node_splits
node_foldtrain -->|"builds datasets"| node_dataset
node_foldtrain -->|"loads segmenter"| node_segmenter
node_dataset -->|"runs segmentation"| node_segmenter
node_dataset -->|"derives targets"| node_heatmaps
node_segmenter -->|"provides masks"| node_heatmaps
node_heatmaps -.->|"uses contour points"| node_points
node_foldtrain -->|"initializes teacher"| node_teacher
node_teacher -->|"runs segmentation"| node_segmenter
node_teacher -->|"generates heatmaps"| node_heatmaps
node_teacher -.->|"extracts masks"| node_yolomasks
node_yolomasks -->|"runs tracking"| node_yolo
node_foldtrain -->|"initializes student"| node_student
node_foldtrain -->|"runs training"| node_trainer
node_trainer -->|"updates weights"| node_student
node_trainer -->|"gets targets"| node_teacher
node_trainer -.->|"computes loss"| node_losses
node_foldtrain -->|"writes metrics"| node_logs
node_evalentry -->|"loads evaluation data"| node_dataset
node_evalentry -->|"loads weights"| node_checkpoints
node_evalentry -->|"uses teacher"| node_teacher
node_evalentry -->|"runs inference"| node_student
node_evalentry -->|"measures distances"| node_analysis
node_evalentry -->|"writes metrics"| node_logs
node_evalentry -.->|"writes optional video"| node_video

click node_foldtrain "https://github.com/oassal/robio2024/blob/main/KD_Kfold_train.py"
click node_evalentry "https://github.com/oassal/robio2024/blob/main/eval.py"
click node_dataset "https://github.com/oassal/robio2024/blob/main/dataset_kd.py"
click node_splits "https://github.com/oassal/robio2024/blob/main/trainers/kfold_functions.py"
click node_teacher "https://github.com/oassal/robio2024/blob/main/models/models_teacher.py"
click node_segmenter "https://github.com/oassal/robio2024/blob/main/models/SegModel.py"
click node_yolomasks "https://github.com/oassal/robio2024/blob/main/utils/yolo_functions.py"
click node_heatmaps "https://github.com/oassal/robio2024/blob/main/utils/utils_new.py"
click node_points "https://github.com/oassal/robio2024/blob/main/utils/normals.py"
click node_student "https://github.com/oassal/robio2024/blob/main/models/models_student.py"
click node_trainer "https://github.com/oassal/robio2024/blob/main/trainers/trainers.py"
click node_losses "https://github.com/oassal/robio2024/blob/main/losses.py"
click node_analysis "https://github.com/oassal/robio2024/blob/main/trainers/heatmap_analysis.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_dataset toneBlue
class node_teacher,node_segmenter,node_yolomasks,node_heatmaps,node_points toneAmber
class node_foldtrain,node_evalentry,node_splits,node_trainer,node_losses toneMint
class node_student,node_analysis toneRose
class node_user,node_rgbdata,node_checkpoints,node_logs,node_video,node_yolo toneIndigo
```

## Components

### Data and targets

- [KD dataset](https://github.com/oassal/robio2024/blob/main/dataset_kd.py): `dataset_kd.py`

### Teacher generation

- [Teacher wrapper](https://github.com/oassal/robio2024/blob/main/models/models_teacher.py): `models/models_teacher.py`
- [DeepLab segmenter](https://github.com/oassal/robio2024/blob/main/models/SegModel.py): `models/SegModel.py`
- [YOLO mask extraction](https://github.com/oassal/robio2024/blob/main/utils/yolo_functions.py): `utils/yolo_functions.py`
- [Heatmap generation](https://github.com/oassal/robio2024/blob/main/utils/utils_new.py): `utils/utils_new.py`
- [Contour geometry](https://github.com/oassal/robio2024/blob/main/utils/normals.py): `utils/normals.py`

### Training and evaluation

- [K-fold training](https://github.com/oassal/robio2024/blob/main/KD_Kfold_train.py): `KD_Kfold_train.py`
- [Evaluation entry](https://github.com/oassal/robio2024/blob/main/eval.py): `eval.py`
- [Fold splits](https://github.com/oassal/robio2024/blob/main/trainers/kfold_functions.py): `trainers/kfold_functions.py`
- [Training step](https://github.com/oassal/robio2024/blob/main/trainers/trainers.py): `trainers/trainers.py`
- [Distillation losses](https://github.com/oassal/robio2024/blob/main/losses.py): `losses.py`

### Model and analysis

- [ResNet VAE](https://github.com/oassal/robio2024/blob/main/models/models_student.py): `models/models_student.py`
- [Heatmap metrics](https://github.com/oassal/robio2024/blob/main/trainers/heatmap_analysis.py): `trainers/heatmap_analysis.py`

### Other components

- **Experimenter**
- **RGB images and labels**
- **Model checkpoints**
- **TensorBoard logs**
- **Evaluation video**
- **YOLO model**

## Connections

- 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