# tue-mps/eomt architecture

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

- Interactive diagram: https://gitdiagram.com/tue-mps/eomt
- Repository: https://github.com/tue-mps/eomt

## Overview

EoMT is a ViT-based image-segmentation research codebase: users configure training or validation, load dataset archives, and obtain instance, panoptic, or semantic predictions and metrics. The graph shows the documented CLI workflow, shared data and model layers, task-specific training branches, and the DINOv2/DINOv3 model choices. Several dataset and training relationships are visible in the source index, but not their full implementations; uncertain links are omitted rather than inferred.

Main entry points: [`datasets`](https://github.com/tue-mps/eomt/tree/master/datasets) (Dataset loaders), [`training/lightning_module.py`](https://github.com/tue-mps/eomt/blob/master/training/lightning_module.py) (Training and evaluation), [`main.py`](https://github.com/tue-mps/eomt/blob/master/main.py) (Train / validate CLI), [`datasets/coco_panoptic.py`](https://github.com/tue-mps/eomt/blob/master/datasets/coco_panoptic.py) (COCO panoptic)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_entry["CLI and configs"]
  node_cli["Train / validate CLI<br/>[main.py]"]
  node_configs["Task configs"]
end

subgraph group_data["Data pipeline"]
  node_datasets["Dataset loaders"]
  node_dataset_core["Archive dataset<br/>[dataset.py]"]
  node_data_module["Batch collation"]
  node_transforms["Image transforms<br/>[transforms.py]"]
  node_coco_panoptic["COCO panoptic<br/>[coco_panoptic.py]"]
  node_ade_panoptic["ADE20K panoptic<br/>[ade20k_panoptic.py]"]
  node_coco_instance["COCO instance<br/>[coco_instance.py]"]
  node_semantic_data["Semantic datasets<br/>[ade20k_semantic.py]"]
  node_cityscapes_data["Cityscapes semantic"]
end

subgraph group_model["Segmentation model"]
  node_model["EoMT network<br/>[eomt.py]"]
  node_encoder["ViT encoder<br/>[vit.py]"]
  node_scale["Mask upscaling<br/>[scale_block.py]"]
end

subgraph group_training["Training and tasks"]
  node_trainer["Training and evaluation"]
  node_instance["Instance task"]
  node_panoptic["Panoptic task"]
  node_semantic["Semantic task"]
  node_loss["Mask classification loss"]
  node_schedule["Warmup polynomial schedule"]
end

subgraph group_models["Model resources"]
  node_notebook["Quick inference<br/>[inference.ipynb]"]
  node_zoo["Pretrained model zoo"]
  node_backbones["DINOv2 / DINOv3"]
end

node_user(("Researcher"))

node_user -->|"starts training or evaluation"| node_cli
node_cli -->|"loads configuration"| node_configs
node_cli -->|"configures data module"| node_datasets
node_cli -->|"configures model task"| node_trainer
node_datasets -.->|"provides dataset option"| node_coco_panoptic
node_datasets -.->|"provides dataset option"| node_ade_panoptic
node_datasets -.->|"provides dataset option"| node_coco_instance
node_datasets -.->|"provides dataset option"| node_semantic_data
node_datasets -.->|"provides dataset option"| node_cityscapes_data
node_coco_panoptic -->|"extends"| node_data_module
node_coco_panoptic -->|"constructs"| node_dataset_core
node_coco_panoptic -->|"uses"| node_transforms
node_ade_panoptic -->|"extends"| node_data_module
node_ade_panoptic -->|"constructs"| node_dataset_core
node_ade_panoptic -->|"uses"| node_transforms
node_coco_instance -->|"uses"| node_dataset_core
node_coco_instance -->|"uses"| node_transforms
node_model -->|"uses"| node_scale
node_model -.->|"encodes image and queries"| node_encoder
node_trainer -->|"uses schedule"| node_schedule
node_instance -->|"extends"| node_trainer
node_instance -->|"computes loss"| node_loss
node_panoptic -->|"extends"| node_trainer
node_panoptic -->|"uses loss"| node_loss
node_semantic -->|"extends"| node_trainer
node_semantic -->|"uses loss"| node_loss
node_user -->|"runs quick inference"| node_notebook
node_user -->|"selects pretrained weights"| node_zoo
node_configs -->|"selects DINO variant"| node_backbones

click node_cli "https://github.com/tue-mps/eomt/blob/master/main.py"
click node_configs "https://github.com/tue-mps/eomt/tree/master/configs"
click node_datasets "https://github.com/tue-mps/eomt/tree/master/datasets"
click node_dataset_core "https://github.com/tue-mps/eomt/blob/master/datasets/dataset.py"
click node_data_module "https://github.com/tue-mps/eomt/blob/master/datasets/lightning_data_module.py"
click node_transforms "https://github.com/tue-mps/eomt/blob/master/datasets/transforms.py"
click node_coco_panoptic "https://github.com/tue-mps/eomt/blob/master/datasets/coco_panoptic.py"
click node_ade_panoptic "https://github.com/tue-mps/eomt/blob/master/datasets/ade20k_panoptic.py"
click node_coco_instance "https://github.com/tue-mps/eomt/blob/master/datasets/coco_instance.py"
click node_semantic_data "https://github.com/tue-mps/eomt/blob/master/datasets/ade20k_semantic.py"
click node_cityscapes_data "https://github.com/tue-mps/eomt/blob/master/datasets/cityscapes_semantic.py"
click node_model "https://github.com/tue-mps/eomt/blob/master/models/eomt.py"
click node_encoder "https://github.com/tue-mps/eomt/blob/master/models/vit.py"
click node_scale "https://github.com/tue-mps/eomt/blob/master/models/scale_block.py"
click node_trainer "https://github.com/tue-mps/eomt/blob/master/training/lightning_module.py"
click node_instance "https://github.com/tue-mps/eomt/blob/master/training/mask_classification_instance.py"
click node_panoptic "https://github.com/tue-mps/eomt/blob/master/training/mask_classification_panoptic.py"
click node_semantic "https://github.com/tue-mps/eomt/blob/master/training/mask_classification_semantic.py"
click node_loss "https://github.com/tue-mps/eomt/blob/master/training/mask_classification_loss.py"
click node_schedule "https://github.com/tue-mps/eomt/blob/master/training/two_stage_warmup_poly_schedule.py"
click node_notebook "https://github.com/tue-mps/eomt/blob/master/inference.ipynb"
click node_zoo "https://github.com/tue-mps/eomt/tree/master/model_zoo"
click node_backbones "https://github.com/tue-mps/eomt/tree/master/configs"

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_cli,node_configs toneBlue
class node_datasets,node_dataset_core,node_data_module,node_transforms,node_coco_panoptic,node_ade_panoptic,node_coco_instance,node_semantic_data,node_cityscapes_data toneAmber
class node_model,node_encoder,node_scale toneMint
class node_trainer,node_instance,node_panoptic,node_semantic,node_loss,node_schedule toneRose
class node_notebook,node_zoo,node_backbones,node_user toneIndigo
```

## Components

### CLI and configs

- [Train / validate CLI](https://github.com/tue-mps/eomt/blob/master/main.py): `main.py`
- [Task configs](https://github.com/tue-mps/eomt/tree/master/configs): `configs`

### Data pipeline

- [Dataset loaders](https://github.com/tue-mps/eomt/tree/master/datasets): `datasets`
- [Archive dataset](https://github.com/tue-mps/eomt/blob/master/datasets/dataset.py): `datasets/dataset.py`
- [Batch collation](https://github.com/tue-mps/eomt/blob/master/datasets/lightning_data_module.py): `datasets/lightning_data_module.py`
- [Image transforms](https://github.com/tue-mps/eomt/blob/master/datasets/transforms.py): `datasets/transforms.py`
- [COCO panoptic](https://github.com/tue-mps/eomt/blob/master/datasets/coco_panoptic.py): `datasets/coco_panoptic.py`
- [ADE20K panoptic](https://github.com/tue-mps/eomt/blob/master/datasets/ade20k_panoptic.py): `datasets/ade20k_panoptic.py`
- [COCO instance](https://github.com/tue-mps/eomt/blob/master/datasets/coco_instance.py): `datasets/coco_instance.py`
- [Semantic datasets](https://github.com/tue-mps/eomt/blob/master/datasets/ade20k_semantic.py): `datasets/ade20k_semantic.py`
- [Cityscapes semantic](https://github.com/tue-mps/eomt/blob/master/datasets/cityscapes_semantic.py): `datasets/cityscapes_semantic.py`

### Segmentation model

- [EoMT network](https://github.com/tue-mps/eomt/blob/master/models/eomt.py): `models/eomt.py`
- [ViT encoder](https://github.com/tue-mps/eomt/blob/master/models/vit.py): `models/vit.py`
- [Mask upscaling](https://github.com/tue-mps/eomt/blob/master/models/scale_block.py): `models/scale_block.py`

### Training and tasks

- [Training and evaluation](https://github.com/tue-mps/eomt/blob/master/training/lightning_module.py): `training/lightning_module.py`
- [Instance task](https://github.com/tue-mps/eomt/blob/master/training/mask_classification_instance.py): `training/mask_classification_instance.py`
- [Panoptic task](https://github.com/tue-mps/eomt/blob/master/training/mask_classification_panoptic.py): `training/mask_classification_panoptic.py`
- [Semantic task](https://github.com/tue-mps/eomt/blob/master/training/mask_classification_semantic.py): `training/mask_classification_semantic.py`
- [Mask classification loss](https://github.com/tue-mps/eomt/blob/master/training/mask_classification_loss.py): `training/mask_classification_loss.py`
- [Warmup polynomial schedule](https://github.com/tue-mps/eomt/blob/master/training/two_stage_warmup_poly_schedule.py): `training/two_stage_warmup_poly_schedule.py`

### Model resources

- [Quick inference](https://github.com/tue-mps/eomt/blob/master/inference.ipynb): `inference.ipynb`
- [Pretrained model zoo](https://github.com/tue-mps/eomt/tree/master/model_zoo): `model_zoo`
- [DINOv2 / DINOv3](https://github.com/tue-mps/eomt/tree/master/configs): `configs`

### Other components

- **Researcher**

## Connections

- Researcher → Train / validate CLI: starts training or evaluation (evidence: `README.md`)
- Train / validate CLI → Task configs: loads configuration (evidence: `README.md`)
- Train / validate CLI → Dataset loaders: configures data module (evidence: `main.py`)
- Train / validate CLI → Training and evaluation: configures model task (evidence: `README.md`)
- Dataset loaders → COCO panoptic: provides dataset option (evidence: `datasets/coco_panoptic.py`)
- Dataset loaders → ADE20K panoptic: provides dataset option (evidence: `datasets/ade20k_panoptic.py`)
- Dataset loaders → COCO instance: provides dataset option (evidence: `datasets/coco_instance.py`)
- Dataset loaders → Semantic datasets: provides dataset option (evidence: `datasets/ade20k_semantic.py`)
- Dataset loaders → Cityscapes semantic: provides dataset option (evidence: `datasets/cityscapes_semantic.py`)
- COCO panoptic → Batch collation: extends (evidence: `datasets/coco_panoptic.py`)
- COCO panoptic → Archive dataset: constructs (evidence: `datasets/coco_panoptic.py`)
- COCO panoptic → Image transforms: uses (evidence: `datasets/coco_panoptic.py`)
- ADE20K panoptic → Batch collation: extends (evidence: `datasets/ade20k_panoptic.py`)
- ADE20K panoptic → Archive dataset: constructs (evidence: `datasets/ade20k_panoptic.py`)
- ADE20K panoptic → Image transforms: uses (evidence: `datasets/ade20k_panoptic.py`)
- COCO instance → Archive dataset: uses (evidence: `datasets/coco_instance.py`)
- COCO instance → Image transforms: uses (evidence: `datasets/coco_instance.py`)
- EoMT network → Mask upscaling: uses (evidence: `models/eomt.py`)
- EoMT network → ViT encoder: encodes image and queries
- Training and evaluation → Warmup polynomial schedule: uses schedule (evidence: `training/lightning_module.py`)
- Instance task → Training and evaluation: extends (evidence: `training/mask_classification_instance.py`)
- Instance task → Mask classification loss: computes loss (evidence: `training/mask_classification_instance.py`)
- Panoptic task → Training and evaluation: extends (evidence: `training/mask_classification_panoptic.py`)
- Panoptic task → Mask classification loss: uses loss (evidence: `training/mask_classification_panoptic.py`)
- Semantic task → Training and evaluation: extends (evidence: `training/mask_classification_semantic.py`)
- Semantic task → Mask classification loss: uses loss (evidence: `training/mask_classification_semantic.py`)
- Researcher → Quick inference: runs quick inference (evidence: `README.md`)
- Researcher → Pretrained model zoo: selects pretrained weights (evidence: `README.md`)
- Task configs → DINOv2 / DINOv3: selects DINO variant (evidence: `README.md`)