# akhasanovv/admm-implementation architecture

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

- Interactive diagram: https://gitdiagram.com/akhasanovv/admm-implementation
- Repository: https://github.com/akhasanovv/admm-implementation

## Overview

#### Purpose and entry points
This is a PyTorch reconstruction system for mask-based lensless images, implementing fixed ADMM, trainable unrolled Le-ADMM, and modular pre/post-processing variants. Primary executables are `train.py`, `inference.py`, `calculate_metrics.py`, and `benchmark_speed.py`; `demo.ipynb` is a Colab-oriented demonstration.

#### Configuration-driven composition
Runtime behavior is assembled from hierarchical YAML configs: model, dataset split, dataloader, transforms, metrics, and experiment writer are independently selectable/overridable from the CLI. Core defaults and inference composition live in `src/configs/baseline.yaml`, `src/configs/inference.yaml`, and model presets under `src/configs/model/`.

#### Data ingestion and physical inputs
Datasets provide lensless measurements, per-image mask/PSF arrays, and optional lensed ground truth. The DigiCam dataset and custom directory layout are handled through `src/datasets/lensless.py`, `src/datasets/base_dataset.py`, and dataset profiles in `src/configs/datasets/`. Batch construction is separated in `src/datasets/collate.py`; lensless-specific preprocessing and optical/sensor representations reside in `src/lensless_helpers/preprocessor.py`, `src/lensless_helpers/psf.py`, and `src/lensless_helpers/sensor.py`.

#### Reconstruction model family
The reconstruction boundary is the model package. Classical fixed-parameter ADMM is in `src/model/admm.py`; learned unrolled ADMM is in `src/model/le_admm.py`; modular learned pre-/post-processing composition is in `src/model/modular_le_admm.py`. The learned denoising/prior network used by these methods is `src/model/drunet.py`, while `src/model/baseline_model.py` supplies the base model contract.

#### Training and checkpoint lifecycle
`train.py` constructs the configured pipeline and delegates optimization lifecycle to `src/trainer/trainer.py`, extending shared control logic in `src/trainer/base_trainer.py`. Training consumes reconstructed outputs and targets through configurable losses, notably `src/loss/admm_loss.py`, and writes model checkpoints under the configured run directory. Fixed ADMM remains usable without a training stage.

#### Inference, persistence, and metrics
`inference.py` uses `src/trainer/inferencer.py` to load optional pretrained weights, execute reconstruction over a split, and save predictions. Metrics are pluggable: reconstruction quality calculations and aggregation are centered in `src/metrics/reconstruction.py`, `src/metrics/base_metric.py`, and `src/metrics/tracker.py`. `calculate_metrics.py` supports post hoc comparison of saved predictions against external ground truth directories.

#### External services and runtime dependencies
The material runtime is Python with PyTorch and its configured package dependencies in `requirements.txt`. Experiment telemetry can target Comet ML or Weights & Biases via `src/logger/cometml.py`, `src/logger/wandb.py`, and `src/logger/logger.py`; Comet requires an API key. Pretrained model artifacts are externally hosted on Hugging Face, while the demo also retrieves data/checkpoints from external sources.

Main entry points: [`src/trainer/trainer.py`](https://github.com/akhasanovv/admm-implementation/blob/main/src/trainer/trainer.py) (Trainer), [`src/configs/baseline.yaml`](https://github.com/akhasanovv/admm-implementation/blob/main/src/configs/baseline.yaml) (YAML composition), [`src/trainer/inferencer.py`](https://github.com/akhasanovv/admm-implementation/blob/main/src/trainer/inferencer.py) (Inferencer), [`src/model/le_admm.py`](https://github.com/akhasanovv/admm-implementation/blob/main/src/model/le_admm.py) (Unrolled Le-ADMM)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_entry["Entrypoints"]
  node_train["train.py<br/>training CLI<br/>[train.py]"]
  node_inference["inference.py<br/>inference CLI<br/>[inference.py]"]
  node_posthoc_metrics["calculate_metrics.py<br/>metrics CLI"]
end

subgraph group_config["Configuration"]
  node_configs["YAML composition<br/>runtime configuration<br/>[baseline.yaml]"]
  node_inference_config["Inference config<br/>runtime configuration<br/>[inference.yaml]"]
end

subgraph group_data["Data &amp; optics"]
  node_dataset["Lensless datasets<br/>dataset layer<br/>[lensless.py]"]
  node_batching["Batch collation<br/>batch builder<br/>[collate.py]"]
  node_optics["Optical preprocessing<br/>physical input processing<br/>[preprocessor.py]"]
end

subgraph group_models["Reconstruction models"]
  node_model_contract["Base model contract<br/>model interface<br/>[baseline_model.py]"]
  node_admm["Fixed ADMM<br/>classical reconstructor<br/>[admm.py]"]
  node_le_admm["Unrolled Le-ADMM<br/>trainable reconstructor<br/>[le_admm.py]"]
  node_modular_admm["Modular Le-ADMM<br/>composed reconstructor<br/>[modular_le_admm.py]"]
  node_drunet["DRUNet prior<br/>learned denoiser<br/>[drunet.py]"]
end

subgraph group_runtime["Training &amp; inference"]
  node_trainer["Trainer<br/>optimization lifecycle<br/>[trainer.py]"]
  node_inferencer["Inferencer<br/>prediction lifecycle<br/>[inferencer.py]"]
  node_admm_loss["ADMM loss<br/>reconstruction loss<br/>[admm_loss.py]"]
end

subgraph group_outputs["Evaluation &amp; services"]
  node_metrics["Reconstruction metrics<br/>metric calculations<br/>[reconstruction.py]"]
  node_predictions[("Saved predictions<br/>reconstruction artifacts")]
  node_checkpoints[("Run checkpoints<br/>model artifacts")]
  node_telemetry["Comet ML / W&amp;B<br/>experiment telemetry<br/>[logger.py]"]
  node_huggingface[("Hugging Face<br/>pretrained artifacts")]
end

node_train -->|"composes"| node_configs
node_inference -->|"composes"| node_inference_config
node_configs -->|"selects split"| node_dataset
node_configs -.->|"selects model"| node_admm
node_configs -.->|"selects model"| node_le_admm
node_configs -.->|"selects model"| node_modular_admm
node_dataset -->|"samples"| node_batching
node_batching -->|"batches inputs"| node_optics
node_optics -->|"measurements and PSF"| node_admm
node_optics -->|"measurements and PSF"| node_le_admm
node_optics -->|"measurements and PSF"| node_modular_admm
node_model_contract -.->|"implements"| node_admm
node_model_contract -.->|"implements"| node_le_admm
node_model_contract -.->|"implements"| node_modular_admm
node_drunet -->|"learned prior"| node_le_admm
node_drunet -->|"learned prior"| node_modular_admm
node_train -->|"delegates"| node_trainer
node_trainer -->|"optimizes"| node_le_admm
node_trainer -->|"optimizes"| node_modular_admm
node_admm_loss -->|"supervises"| node_trainer
node_trainer -->|"writes"| node_checkpoints
node_trainer -.->|"logs to"| node_telemetry
node_inference -->|"delegates"| node_inferencer
node_checkpoints -->|"optional weights"| node_inferencer
node_huggingface -.->|"pretrained weights"| node_inferencer
node_inferencer -->|"saves"| node_predictions
node_predictions -->|"compares"| node_posthoc_metrics
node_posthoc_metrics -->|"calculates"| node_metrics
node_inferencer -.->|"evaluates"| node_metrics

click node_train "https://github.com/akhasanovv/admm-implementation/blob/main/train.py"
click node_inference "https://github.com/akhasanovv/admm-implementation/blob/main/inference.py"
click node_posthoc_metrics "https://github.com/akhasanovv/admm-implementation/blob/main/calculate_metrics.py"
click node_configs "https://github.com/akhasanovv/admm-implementation/blob/main/src/configs/baseline.yaml"
click node_inference_config "https://github.com/akhasanovv/admm-implementation/blob/main/src/configs/inference.yaml"
click node_dataset "https://github.com/akhasanovv/admm-implementation/blob/main/src/datasets/lensless.py"
click node_batching "https://github.com/akhasanovv/admm-implementation/blob/main/src/datasets/collate.py"
click node_optics "https://github.com/akhasanovv/admm-implementation/blob/main/src/lensless_helpers/preprocessor.py"
click node_model_contract "https://github.com/akhasanovv/admm-implementation/blob/main/src/model/baseline_model.py"
click node_admm "https://github.com/akhasanovv/admm-implementation/blob/main/src/model/admm.py"
click node_le_admm "https://github.com/akhasanovv/admm-implementation/blob/main/src/model/le_admm.py"
click node_modular_admm "https://github.com/akhasanovv/admm-implementation/blob/main/src/model/modular_le_admm.py"
click node_drunet "https://github.com/akhasanovv/admm-implementation/blob/main/src/model/drunet.py"
click node_trainer "https://github.com/akhasanovv/admm-implementation/blob/main/src/trainer/trainer.py"
click node_inferencer "https://github.com/akhasanovv/admm-implementation/blob/main/src/trainer/inferencer.py"
click node_admm_loss "https://github.com/akhasanovv/admm-implementation/blob/main/src/loss/admm_loss.py"
click node_metrics "https://github.com/akhasanovv/admm-implementation/blob/main/src/metrics/reconstruction.py"
click node_telemetry "https://github.com/akhasanovv/admm-implementation/blob/main/src/logger/logger.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_train,node_inference,node_posthoc_metrics toneBlue
class node_configs,node_inference_config toneAmber
class node_dataset,node_batching,node_optics toneMint
class node_model_contract,node_admm,node_le_admm,node_modular_admm,node_drunet toneRose
class node_trainer,node_inferencer,node_admm_loss toneIndigo
class node_metrics,node_predictions,node_checkpoints,node_telemetry,node_huggingface toneTeal
```

## Components

### Entrypoints

- [train.py](https://github.com/akhasanovv/admm-implementation/blob/main/train.py): training CLI · `train.py`
- [inference.py](https://github.com/akhasanovv/admm-implementation/blob/main/inference.py): inference CLI · `inference.py`
- [calculate_metrics.py](https://github.com/akhasanovv/admm-implementation/blob/main/calculate_metrics.py): metrics CLI · Compares saved predictions with external ground truth. · `calculate_metrics.py`

### Configuration

- [YAML composition](https://github.com/akhasanovv/admm-implementation/blob/main/src/configs/baseline.yaml): runtime configuration · CLI-overridable model, data, transforms, metrics, and writer selection. · `src/configs/baseline.yaml`
- [Inference config](https://github.com/akhasanovv/admm-implementation/blob/main/src/configs/inference.yaml): runtime configuration · `src/configs/inference.yaml`

### Data & optics

- [Lensless datasets](https://github.com/akhasanovv/admm-implementation/blob/main/src/datasets/lensless.py): dataset layer · Loads measurements, per-image mask/PSF, and optional lensed targets. · `src/datasets/lensless.py`
- [Batch collation](https://github.com/akhasanovv/admm-implementation/blob/main/src/datasets/collate.py): batch builder · `src/datasets/collate.py`
- [Optical preprocessing](https://github.com/akhasanovv/admm-implementation/blob/main/src/lensless_helpers/preprocessor.py): physical input processing · Prepares lensless inputs and PSF/sensor representations. · `src/lensless_helpers/preprocessor.py`

### Reconstruction models

- [Base model contract](https://github.com/akhasanovv/admm-implementation/blob/main/src/model/baseline_model.py): model interface · `src/model/baseline_model.py`
- [Fixed ADMM](https://github.com/akhasanovv/admm-implementation/blob/main/src/model/admm.py): classical reconstructor · Runs fixed-parameter reconstruction without training. · `src/model/admm.py`
- [Unrolled Le-ADMM](https://github.com/akhasanovv/admm-implementation/blob/main/src/model/le_admm.py): trainable reconstructor · `src/model/le_admm.py`
- [Modular Le-ADMM](https://github.com/akhasanovv/admm-implementation/blob/main/src/model/modular_le_admm.py): composed reconstructor · Adds learned pre- and post-processing around unrolled ADMM. · `src/model/modular_le_admm.py`
- [DRUNet prior](https://github.com/akhasanovv/admm-implementation/blob/main/src/model/drunet.py): learned denoiser · `src/model/drunet.py`

### Training & inference

- [Trainer](https://github.com/akhasanovv/admm-implementation/blob/main/src/trainer/trainer.py): optimization lifecycle · Optimizes configured trainable reconstruction models and writes checkpoints. · `src/trainer/trainer.py`
- [Inferencer](https://github.com/akhasanovv/admm-implementation/blob/main/src/trainer/inferencer.py): prediction lifecycle · Loads optional weights, reconstructs a split, and saves predictions. · `src/trainer/inferencer.py`
- [ADMM loss](https://github.com/akhasanovv/admm-implementation/blob/main/src/loss/admm_loss.py): reconstruction loss · `src/loss/admm_loss.py`

### Evaluation & services

- [Reconstruction metrics](https://github.com/akhasanovv/admm-implementation/blob/main/src/metrics/reconstruction.py): metric calculations · Computes and aggregates reconstruction quality. · `src/metrics/reconstruction.py`
- **Saved predictions**: reconstruction artifacts
- **Run checkpoints**: model artifacts
- [Comet ML / W&B](https://github.com/akhasanovv/admm-implementation/blob/main/src/logger/logger.py): experiment telemetry · Optional external experiment logging. · `src/logger/logger.py`
- **Hugging Face**: pretrained artifacts · External host for pretrained model artifacts.

## Connections

- train.py → YAML composition: composes
- inference.py → Inference config: composes
- YAML composition → Lensless datasets: selects split
- YAML composition → Fixed ADMM: selects model
- YAML composition → Unrolled Le-ADMM: selects model
- YAML composition → Modular Le-ADMM: selects model
- Lensless datasets → Batch collation: samples
- Batch collation → Optical preprocessing: batches inputs
- Optical preprocessing → Fixed ADMM: measurements and PSF
- Optical preprocessing → Unrolled Le-ADMM: measurements and PSF
- Optical preprocessing → Modular Le-ADMM: measurements and PSF
- Base model contract → Fixed ADMM: implements
- Base model contract → Unrolled Le-ADMM: implements
- Base model contract → Modular Le-ADMM: implements
- DRUNet prior → Unrolled Le-ADMM: learned prior
- DRUNet prior → Modular Le-ADMM: learned prior
- train.py → Trainer: delegates
- Trainer → Unrolled Le-ADMM: optimizes
- Trainer → Modular Le-ADMM: optimizes
- ADMM loss → Trainer: supervises
- Trainer → Run checkpoints: writes
- Trainer → Comet ML / W&B: logs to
- inference.py → Inferencer: delegates
- Run checkpoints → Inferencer: optional weights
- Hugging Face → Inferencer: pretrained weights
- Inferencer → Saved predictions: saves
- Saved predictions → calculate_metrics.py: compares
- calculate_metrics.py → Reconstruction metrics: calculates
- Inferencer → Reconstruction metrics: evaluates