# m1dsolo/echocardmae architecture

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

- Interactive diagram: https://gitdiagram.com/m1dsolo/echocardmae
- Repository: https://github.com/m1dsolo/echocardmae

## Overview

EchoCardMAE is a research workflow for echocardiography video representation learning, ejection-fraction (EF) prediction, and cardiac segmentation. Researchers prepare dataset-specific arrays, optionally pretrain a masked video autoencoder, then fine-tune or evaluate task models. The graph retains the separate data-preparation and task branches, plus the shared transformer internals. Some preprocessing and model files were not excerpted; their workflow placement is based on the README, and unsampled implementation details are omitted.

Main entry points: [`utils.py`](https://github.com/m1dsolo/echocardmae/blob/main/utils.py) (Array dataset utilities), [`pretrain.py`](https://github.com/m1dsolo/echocardmae/blob/main/pretrain.py) (MAE pretraining), [`echonet/train_seg.py`](https://github.com/m1dsolo/echocardmae/blob/main/echonet/train_seg.py) (Segmentation training), [`echonet/train_ef.py`](https://github.com/m1dsolo/echocardmae/blob/main/echonet/train_ef.py) (EF fine-tuning)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Data preparation"]
  node_efprep["EF preprocessing<br/>[avi2npy.py]"]
  node_segprep["Segmentation preprocessing<br/>[avi2edes_npy.py]"]
  node_datasetutil["Array dataset utilities<br/>[utils.py]"]
end

subgraph group_pretraining["Self-supervised pretraining"]
  node_mae["Masked video autoencoder<br/>[echocardmae_vit.py]"]
  node_pretrain["MAE pretraining<br/>[pretrain.py]"]
  node_infonce["Contrastive loss<br/>[infonce_loss.py]"]
  node_transformerutils["Transformer building blocks<br/>[utils.py]"]
end

subgraph group_ef["EF prediction"]
  node_vit["Video transformer<br/>[vit.py]"]
  node_eftrain["EF fine-tuning<br/>[train_ef.py]"]
  node_efval["EF evaluation<br/>[val_ef.py]"]
end

subgraph group_segmentation["Segmentation"]
  node_vitseg["Segmentation transformer<br/>[vit_seg.py]"]
  node_segtrain["Segmentation training<br/>[train_seg.py]"]
  node_diceloss["Dice loss<br/>[dice_loss.py]"]
end

node_researcher(("Researcher"))
node_datasets[("Echo datasets")]

node_researcher -->|"obtains data"| node_datasets
node_datasets -->|"provides videos"| node_efprep
node_datasets -->|"provides videos"| node_segprep
node_researcher -->|"runs preprocessing"| node_efprep
node_researcher -->|"runs preprocessing"| node_segprep
node_efprep -->|"uses utilities"| node_datasetutil
node_segprep -->|"uses utilities"| node_datasetutil
node_researcher -->|"starts training"| node_pretrain
node_pretrain -->|"loads dataset"| node_datasetutil
node_pretrain -->|"trains model"| node_mae
node_pretrain -->|"computes alignment"| node_infonce
node_mae -->|"uses blocks"| node_transformerutils
node_researcher -->|"starts fine-tuning"| node_eftrain
node_eftrain -->|"loads dataset"| node_datasetutil
node_eftrain -->|"trains model"| node_vit
node_vit -->|"uses blocks"| node_transformerutils
node_researcher -->|"runs evaluation"| node_efval
node_efval -->|"loads dataset"| node_datasetutil
node_efval -->|"evaluates model"| node_vit
node_researcher -->|"starts training"| node_segtrain
node_segtrain -->|"loads dataset"| node_datasetutil
node_segtrain -->|"trains model"| node_vitseg
node_segtrain -->|"computes loss"| node_diceloss
node_vitseg -->|"uses blocks"| node_transformerutils

click node_efprep "https://github.com/m1dsolo/echocardmae/blob/main/echonet/avi2npy.py"
click node_segprep "https://github.com/m1dsolo/echocardmae/blob/main/echonet/avi2edes_npy.py"
click node_datasetutil "https://github.com/m1dsolo/echocardmae/blob/main/utils.py"
click node_mae "https://github.com/m1dsolo/echocardmae/blob/main/model/echocardmae_vit.py"
click node_pretrain "https://github.com/m1dsolo/echocardmae/blob/main/pretrain.py"
click node_infonce "https://github.com/m1dsolo/echocardmae/blob/main/loss/infonce_loss.py"
click node_vit "https://github.com/m1dsolo/echocardmae/blob/main/model/vit.py"
click node_eftrain "https://github.com/m1dsolo/echocardmae/blob/main/echonet/train_ef.py"
click node_efval "https://github.com/m1dsolo/echocardmae/blob/main/echonet/val_ef.py"
click node_vitseg "https://github.com/m1dsolo/echocardmae/blob/main/model/vit_seg.py"
click node_segtrain "https://github.com/m1dsolo/echocardmae/blob/main/echonet/train_seg.py"
click node_diceloss "https://github.com/m1dsolo/echocardmae/blob/main/loss/dice_loss.py"
click node_transformerutils "https://github.com/m1dsolo/echocardmae/blob/main/model/utils.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_efprep,node_segprep,node_datasetutil toneBlue
class node_mae,node_pretrain,node_infonce,node_transformerutils,node_datasets toneAmber
class node_vit,node_eftrain,node_efval toneMint
class node_vitseg,node_segtrain,node_diceloss toneRose
class node_researcher toneIndigo
```

## Components

### Data preparation

- [EF preprocessing](https://github.com/m1dsolo/echocardmae/blob/main/echonet/avi2npy.py): `echonet/avi2npy.py`
- [Segmentation preprocessing](https://github.com/m1dsolo/echocardmae/blob/main/echonet/avi2edes_npy.py): `echonet/avi2edes_npy.py`
- [Array dataset utilities](https://github.com/m1dsolo/echocardmae/blob/main/utils.py): `utils.py`

### Self-supervised pretraining

- [Masked video autoencoder](https://github.com/m1dsolo/echocardmae/blob/main/model/echocardmae_vit.py): `model/echocardmae_vit.py`
- [MAE pretraining](https://github.com/m1dsolo/echocardmae/blob/main/pretrain.py): `pretrain.py`
- [Contrastive loss](https://github.com/m1dsolo/echocardmae/blob/main/loss/infonce_loss.py): `loss/infonce_loss.py`
- [Transformer building blocks](https://github.com/m1dsolo/echocardmae/blob/main/model/utils.py): `model/utils.py`

### EF prediction

- [Video transformer](https://github.com/m1dsolo/echocardmae/blob/main/model/vit.py): `model/vit.py`
- [EF fine-tuning](https://github.com/m1dsolo/echocardmae/blob/main/echonet/train_ef.py): `echonet/train_ef.py`
- [EF evaluation](https://github.com/m1dsolo/echocardmae/blob/main/echonet/val_ef.py): `echonet/val_ef.py`

### Segmentation

- [Segmentation transformer](https://github.com/m1dsolo/echocardmae/blob/main/model/vit_seg.py): `model/vit_seg.py`
- [Segmentation training](https://github.com/m1dsolo/echocardmae/blob/main/echonet/train_seg.py): `echonet/train_seg.py`
- [Dice loss](https://github.com/m1dsolo/echocardmae/blob/main/loss/dice_loss.py): `loss/dice_loss.py`

### Other components

- **Researcher**
- **Echo datasets**

## Connections

- Researcher → Echo datasets: obtains data (evidence: `README.md`)
- Echo datasets → EF preprocessing: provides videos (evidence: `README.md`)
- Echo datasets → Segmentation preprocessing: provides videos (evidence: `README.md`)
- Researcher → EF preprocessing: runs preprocessing (evidence: `README.md`)
- Researcher → Segmentation preprocessing: runs preprocessing (evidence: `README.md`)
- EF preprocessing → Array dataset utilities: uses utilities (evidence: `echonet/avi2npy.py`)
- Segmentation preprocessing → Array dataset utilities: uses utilities (evidence: `echonet/avi2edes_npy.py`)
- Researcher → MAE pretraining: starts training (evidence: `README.md`)
- MAE pretraining → Array dataset utilities: loads dataset (evidence: `pretrain.py`)
- MAE pretraining → Masked video autoencoder: trains model (evidence: `pretrain.py`)
- MAE pretraining → Contrastive loss: computes alignment (evidence: `pretrain.py`)
- Masked video autoencoder → Transformer building blocks: uses blocks (evidence: `model/echocardmae_vit.py`)
- Researcher → EF fine-tuning: starts fine-tuning (evidence: `README.md`)
- EF fine-tuning → Array dataset utilities: loads dataset (evidence: `echonet/train_ef.py`)
- EF fine-tuning → Video transformer: trains model (evidence: `echonet/train_ef.py`)
- Video transformer → Transformer building blocks: uses blocks (evidence: `model/vit.py`)
- Researcher → EF evaluation: runs evaluation (evidence: `README.md`)
- EF evaluation → Array dataset utilities: loads dataset (evidence: `echonet/val_ef.py`)
- EF evaluation → Video transformer: evaluates model (evidence: `echonet/val_ef.py`)
- Researcher → Segmentation training: starts training (evidence: `README.md`)
- Segmentation training → Array dataset utilities: loads dataset (evidence: `echonet/train_seg.py`)
- Segmentation training → Segmentation transformer: trains model (evidence: `echonet/train_seg.py`)
- Segmentation training → Dice loss: computes loss (evidence: `echonet/train_seg.py`)
- Segmentation transformer → Transformer building blocks: uses blocks (evidence: `model/vit_seg.py`)