# mdivyanshu97/discovr architecture

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

- Interactive diagram: https://gitdiagram.com/mdivyanshu97/discovr
- Repository: https://github.com/mdivyanshu97/discovr

## Overview

1. **Overall shape**
   - This repo is a research training stack for DISCOVR, a self-supervised echocardiography video representation learner.
   - The main architecture is a **pretraining pipeline** plus a **downstream fine-tuning/evaluation pipeline**, both centered on video clips and checkpointed PyTorch models.

2. **Top-level execution entry points**
   - `scripts/run_mae_pretraining.py` is the main pretraining launcher.
   - `scripts/run_class_finetuning.py` is the downstream supervised training entry point.
   - `run_discovr_encoder.py` is a lightweight inference wrapper for loading a pretrained checkpoint and extracting embeddings.
   - `config/run_pretraining.sh` appears to be a convenience shell launcher for training jobs.

3. **Core model layer**
   - `models/modeling_pretrain.py` and `models/modeling_finetune.py` are the central model definitions for the two training stages.
   - Supporting model components include `models/masking_generator.py`, `models/attn_mask.py`, and `models/info_nce.py`, which suggest the pretraining objective combines masked video modeling, attention/masking logic, and contrastive/distillation-style losses.
   - The README describes a dual-branch design: a **video branch** for temporal dynamics and an **online image branch** for spatial semantics, with cluster distillation between them.

4. **Training orchestration**
   - `engine/engine_for_pretraining.py` contains the epoch/step logic for self-supervised learning, likely handling forward passes, loss computation, optimization, and checkpointing.
   - `engine/engine_for_finetuning.py` handles supervised downstream training and validation.
   - These engines form the boundary between reusable model code and task-specific training behavior.

5. **Data ingestion boundary**
   - `data/datasets/datasets.py` and `data/datasets/kinetics.py` define dataset wrappers and likely CSV/file-based sample indexing.
   - The README and scripts indicate training is driven by paths to videos plus CSV splits for train/val/test.
   - The data layer is where raw echocardiography videos enter the system before any masking or augmentation is applied.

6. **Video preprocessing and augmentation**
   - `data/transforms/` contains the augmentation stack for spatiotemporal inputs: `video_transforms.py`, `volume_transforms.py`, `rand_augment.py`, `mixup.py`, `random_erasing.py`, and helper `functional.py`.
   - This is the main “input shaping” layer, converting decoded clips into model-ready tensors and applying regularization augmentations.
   - `utils/visualize_masks.py` and `scripts/visualise_dataset.py` support debugging/inspection of what the model sees.

7. **Optimization, metrics, and utility layer**
   - `utils/optim_factory.py` centralizes optimizer construction and parameter grouping.
   - `utils/metrics_utils.py` and `utils/utils.py` provide common training utilities, metric tracking, logging helpers, and general runtime support.
   - `utils/utils_datasets.py` likely bridges dataset metadata and training code.

8. **Dataflow through pretraining**
   - The pretraining path is: video files + CSV splits → dataset loaders → transforms/masking → pretraining engine → `modeling_pretrain` → losses such as InfoNCE/distillation/masked prediction → checkpoints.
   - The pretrained checkpoint is then reused either for inference (`run_discovr_encoder.py`) or as initialization for downstream evaluation.

9. **Dataflow through fine-tuning**
   - The fine-tuning path is: labeled video data → dataset/transforms → `engine_for_finetuning.py` → `modeling_finetune.py` → task-specific predictions and metrics.
   - This separation keeps the self-supervised representation learning code distinct from supervised evaluation code.

10. **Runtime and infrastructure considerations**
   - The project is designed for **distributed GPU training** via `torch.distributed.launch` and PyTorch.
   - The README makes FFmpeg/TorchCodec part of the runtime because video decoding is a first-class dependency of the pipeline.
   - Hugging Face is used for checkpoint distribution, but not as a core runtime service.

11. **Standalone TubeViT subtree**
   - `models/TubeViT/` is a self-contained subproject with its own `setup.py`, scripts, dataset/model code, and requirements.
   - It includes training, evaluation, inference, visualization, and ViT weight conversion utilities.
   - Architecturally, this looks like a bundled external model implementation or experimental dependency rather than the main DISCOVR training path.

12. **What an architecture diagram should emphasize**
   - Show three boundaries clearly: **data decoding/augmentation**, **training engines**, and **model/loss definitions**.
   - Show two primary flows: **self-supervised pretraining** and **supervised fine-tuning/inference**.
   - Include the optional/secondary **TubeViT submodule** as a separate model stack, and note that pretrained checkpoints are externalized through Hugging Face.

Main entry points: [`engine/engine_for_pretraining.py`](https://github.com/mdivyanshu97/discovr/blob/master/engine/engine_for_pretraining.py) (Pretrain engine), [`data/transforms/transforms.py`](https://github.com/mdivyanshu97/discovr/blob/master/data/transforms/transforms.py) (Video transforms), [`engine/engine_for_finetuning.py`](https://github.com/mdivyanshu97/discovr/blob/master/engine/engine_for_finetuning.py) (Finetune engine), [`data/datasets/datasets.py`](https://github.com/mdivyanshu97/discovr/blob/master/data/datasets/datasets.py) (Datasets)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_main["DISCOVR stack"]
  node_pretrain_script["Pretrain run<br/>entrypoint"]
  node_finetune_script["Finetune run<br/>entrypoint"]
  node_encoder_script["Encoder infer<br/>entrypoint"]
  node_pretrain_engine["Pretrain engine<br/>training loop"]
  node_finetune_engine["Finetune engine<br/>training loop"]
  node_pretrain_model{{"Pretrain model<br/>dual-branch model"}}
  node_finetune_model{{"Finetune model<br/>classifier model"}}
  node_masking(("Masking"))
  node_attn_mask(("Attn mask<br/>attention utility<br/>[attn_mask.py]"))
  node_info_nce(("InfoNCE<br/>contrastive loss<br/>[info_nce.py]"))
  node_datasets[("Datasets<br/>video dataset<br/>[datasets.py]")]
  node_kinetics[("Kinetics IO<br/>dataset adapter<br/>[kinetics.py]")]
  node_transforms["Video transforms<br/>augmentation pipeline<br/>[transforms.py]"]
  node_video_transforms["Video aug<br/>transform set"]
  node_volume_transforms["Volume aug<br/>transform set"]
  node_optim(("Optimizer<br/>optim factory<br/>[optim_factory.py]"))
  node_metrics(("Metrics<br/>training utils<br/>[metrics_utils.py]"))
end

subgraph group_tubevit["TubeViT submodule"]
  node_tubevit_model{{"TubeViT model<br/>vit stack<br/>[model.py]"}}
  node_tubevit_train["TubeViT scripts<br/>standalone tools<br/>[train.py]"]
end

node_pretrain_script -->|"launches"| node_pretrain_engine
node_finetune_script -->|"launches"| node_finetune_engine
node_encoder_script -->|"loads"| node_pretrain_model
node_pretrain_engine -->|"trains"| node_pretrain_model
node_finetune_engine -->|"trains"| node_finetune_model
node_pretrain_engine -->|"applies"| node_masking
node_pretrain_engine -->|"uses"| node_attn_mask
node_pretrain_engine -->|"optimizes"| node_info_nce
node_datasets -->|"specializes"| node_kinetics
node_pretrain_script -->|"reads"| node_datasets
node_finetune_script -->|"reads"| node_datasets
node_datasets -->|"feeds"| node_transforms
node_transforms -->|"includes"| node_video_transforms
node_transforms -->|"includes"| node_volume_transforms
node_transforms -->|"prepares"| node_pretrain_engine
node_transforms -->|"prepares"| node_finetune_engine
node_pretrain_engine -->|"configures"| node_optim
node_finetune_engine -->|"reports"| node_metrics
node_tubevit_train -->|"uses"| node_tubevit_model
node_tubevit_model -->|"supports"| node_tubevit_train

click node_pretrain_script "https://github.com/mdivyanshu97/discovr/blob/master/scripts/run_mae_pretraining.py"
click node_finetune_script "https://github.com/mdivyanshu97/discovr/blob/master/scripts/run_class_finetuning.py"
click node_encoder_script "https://github.com/mdivyanshu97/discovr/blob/master/run_discovr_encoder.py"
click node_pretrain_engine "https://github.com/mdivyanshu97/discovr/blob/master/engine/engine_for_pretraining.py"
click node_finetune_engine "https://github.com/mdivyanshu97/discovr/blob/master/engine/engine_for_finetuning.py"
click node_pretrain_model "https://github.com/mdivyanshu97/discovr/blob/master/models/modeling_pretrain.py"
click node_finetune_model "https://github.com/mdivyanshu97/discovr/blob/master/models/modeling_finetune.py"
click node_masking "https://github.com/mdivyanshu97/discovr/blob/master/models/masking_generator.py"
click node_attn_mask "https://github.com/mdivyanshu97/discovr/blob/master/models/attn_mask.py"
click node_info_nce "https://github.com/mdivyanshu97/discovr/blob/master/models/info_nce.py"
click node_datasets "https://github.com/mdivyanshu97/discovr/blob/master/data/datasets/datasets.py"
click node_kinetics "https://github.com/mdivyanshu97/discovr/blob/master/data/datasets/kinetics.py"
click node_transforms "https://github.com/mdivyanshu97/discovr/blob/master/data/transforms/transforms.py"
click node_video_transforms "https://github.com/mdivyanshu97/discovr/blob/master/data/transforms/video_transforms.py"
click node_volume_transforms "https://github.com/mdivyanshu97/discovr/blob/master/data/transforms/volume_transforms.py"
click node_optim "https://github.com/mdivyanshu97/discovr/blob/master/utils/optim_factory.py"
click node_metrics "https://github.com/mdivyanshu97/discovr/blob/master/utils/metrics_utils.py"
click node_tubevit_model "https://github.com/mdivyanshu97/discovr/blob/master/models/TubeViT/tubevit/model.py"
click node_tubevit_train "https://github.com/mdivyanshu97/discovr/blob/master/models/TubeViT/scripts/train.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_pretrain_script,node_finetune_script,node_encoder_script,node_pretrain_engine,node_finetune_engine,node_pretrain_model,node_finetune_model,node_masking,node_attn_mask,node_info_nce,node_datasets,node_kinetics,node_transforms,node_video_transforms,node_volume_transforms,node_optim,node_metrics toneBlue
class node_tubevit_model,node_tubevit_train toneAmber
```

## Components

### DISCOVR stack

Primary pretraining, fine-tuning, and inference pipeline for self-supervised echocardiography video representation learning.

- [Pretrain run](https://github.com/mdivyanshu97/discovr/blob/master/scripts/run_mae_pretraining.py): entrypoint · Main launcher for masked self-supervised pretraining. · `scripts/run_mae_pretraining.py`
- [Finetune run](https://github.com/mdivyanshu97/discovr/blob/master/scripts/run_class_finetuning.py): entrypoint · Supervised downstream training entrypoint. · `scripts/run_class_finetuning.py`
- [Encoder infer](https://github.com/mdivyanshu97/discovr/blob/master/run_discovr_encoder.py): entrypoint · Loads a pretrained checkpoint and extracts embeddings. · `run_discovr_encoder.py`
- [Pretrain engine](https://github.com/mdivyanshu97/discovr/blob/master/engine/engine_for_pretraining.py): training loop · Epoch and step orchestration for self-supervised training. · `engine/engine_for_pretraining.py`
- [Finetune engine](https://github.com/mdivyanshu97/discovr/blob/master/engine/engine_for_finetuning.py): training loop · Supervised training and validation orchestration. · `engine/engine_for_finetuning.py`
- [Pretrain model](https://github.com/mdivyanshu97/discovr/blob/master/models/modeling_pretrain.py): dual-branch model · Masked video representation learner with online image/video branches. · `models/modeling_pretrain.py`
- [Finetune model](https://github.com/mdivyanshu97/discovr/blob/master/models/modeling_finetune.py): classifier model · Task-specific downstream prediction model. · `models/modeling_finetune.py`
- [Masking](https://github.com/mdivyanshu97/discovr/blob/master/models/masking_generator.py): masking utils · Mask generation and attention masking logic for pretraining. · `models/masking_generator.py`
- [Attn mask](https://github.com/mdivyanshu97/discovr/blob/master/models/attn_mask.py): attention utility · Attention mask construction for masked sequence modeling. · `models/attn_mask.py`
- [InfoNCE](https://github.com/mdivyanshu97/discovr/blob/master/models/info_nce.py): contrastive loss · Contrastive/distillation-style objective used in training. · `models/info_nce.py`
- [Datasets](https://github.com/mdivyanshu97/discovr/blob/master/data/datasets/datasets.py): video dataset · Dataset wrappers and sample indexing for video splits. · `data/datasets/datasets.py`
- [Kinetics IO](https://github.com/mdivyanshu97/discovr/blob/master/data/datasets/kinetics.py): dataset adapter · Dataset-specific loader and video sample interface. · `data/datasets/kinetics.py`
- [Video transforms](https://github.com/mdivyanshu97/discovr/blob/master/data/transforms/transforms.py): augmentation pipeline · Spatiotemporal decoding, augmentation, mixup, and erasing pipeline. · `data/transforms/transforms.py`
- [Video aug](https://github.com/mdivyanshu97/discovr/blob/master/data/transforms/video_transforms.py): transform set · Clip-level video transformations. · `data/transforms/video_transforms.py`
- [Volume aug](https://github.com/mdivyanshu97/discovr/blob/master/data/transforms/volume_transforms.py): transform set · 3D volume and temporal tensor transforms. · `data/transforms/volume_transforms.py`
- [Optimizer](https://github.com/mdivyanshu97/discovr/blob/master/utils/optim_factory.py): optim factory · Optimizer construction and parameter grouping. · `utils/optim_factory.py`
- [Metrics](https://github.com/mdivyanshu97/discovr/blob/master/utils/metrics_utils.py): training utils · Metric tracking, logging, and runtime helpers. · `utils/metrics_utils.py`

### TubeViT submodule

Bundled standalone vision transformer subtree with its own scripts, model, and dataset code.

- [TubeViT model](https://github.com/mdivyanshu97/discovr/blob/master/models/TubeViT/tubevit/model.py): vit stack · Standalone TubeViT implementation and positional encoding code. · `models/TubeViT/tubevit/model.py`
- [TubeViT scripts](https://github.com/mdivyanshu97/discovr/blob/master/models/TubeViT/scripts/train.py): standalone tools · Separate training, evaluation, inference, and visualization entrypoints. · `models/TubeViT/scripts/train.py`

## Connections

- Pretrain run → Pretrain engine: launches · Starts self-supervised training.
- Finetune run → Finetune engine: launches · Starts downstream supervised training.
- Encoder infer → Pretrain model: loads · Uses pretrained checkpoint for embedding extraction.
- Pretrain engine → Pretrain model: trains · Runs masked representation learning forward/backward passes.
- Finetune engine → Finetune model: trains · Optimizes downstream predictions.
- Pretrain engine → Masking: applies · Generates masked inputs for pretraining.
- Pretrain engine → Attn mask: uses · Builds attention masks for masked modeling.
- Pretrain engine → InfoNCE: optimizes · Computes representation and distillation objectives.
- Datasets → Kinetics IO: specializes · Dataset-specific video loading adapter.
- Pretrain run → Datasets: reads · Creates training and validation loaders from split metadata.
- Finetune run → Datasets: reads · Builds labeled downstream data loaders.
- Datasets → Video transforms: feeds · Sends decoded clips into augmentation pipeline.
- Video transforms → Video aug: includes · Clip-level augmentation stage.
- Video transforms → Volume aug: includes · Tensor/volume transformation stage.
- Video transforms → Pretrain engine: prepares · Produces model-ready masked clips.
- Video transforms → Finetune engine: prepares · Produces model-ready labeled clips.
- Pretrain engine → Optimizer: configures · Constructs optimizer and parameter groups.
- Finetune engine → Metrics: reports · Tracks losses and evaluation metrics.
- TubeViT scripts → TubeViT model: uses · Standalone TubeViT training and inference scripts.
- TubeViT model → TubeViT scripts: supports · Shared model code for the submodule tools.