# ewrfcas/zits-plusplus architecture

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

- Interactive diagram: https://gitdiagram.com/ewrfcas/zits-plusplus
- Repository: https://github.com/ewrfcas/zits-plusplus

## Overview

ZITS++ is an image-inpainting system: users provide images and masks, the runtime prepares inputs and structural guidance, then a learned generator reconstructs masked regions for output. The README documents 256/512 inference, optional object removal, pretrained models, and training-related code. Sampled dataset, trainer, generator, and line-detector code supports the main stages; the inference entry point and some model wiring are absent from the supplied tree or excerpts, so uncertain connections are omitted.

Main entry points: [`trainers/pl_trainers.py`](https://github.com/ewrfcas/zits-plusplus/blob/main/trainers/pl_trainers.py) (Inpainting trainer), [`dataset/dataloader.py`](https://github.com/ewrfcas/zits-plusplus/blob/main/dataset/dataloader.py) (Test dataset), [`base/dataset.py`](https://github.com/ewrfcas/zits-plusplus/blob/main/base/dataset.py) (Training datasets), [`trainers/lsm_hawp/detector.py`](https://github.com/ewrfcas/zits-plusplus/blob/main/trainers/lsm_hawp/detector.py) (LSM-HAWP detector)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_input["Input preparation"]
  node_testdata["Test dataset<br/>[dataloader.py]"]
  node_traindata["Training datasets<br/>[dataset.py]"]
end

subgraph group_structure["Structural priors"]
  node_line_detector["LSM-HAWP detector<br/>[detector.py]"]
  node_segmentation["ADE20K segmentation<br/>[base.py]"]
  node_structure_net["Structure upsampler<br/>[tsr.py]"]
  node_prior_net["Prior networks<br/>[pcp.py]"]
end

subgraph group_models["Inpainting models"]
  node_generator["MAT generator<br/>[mat.py]"]
  node_generator_family["Generator components<br/>[generators.py]"]
  node_discriminator["Discriminator<br/>[discriminators.py]"]
  node_losses["Training losses<br/>[losses.py]"]
end

subgraph group_runtime["Runtime and training"]
  node_trainer["Inpainting trainer<br/>[pl_trainers.py]"]
  node_metric["Inpainting metrics"]
  node_logger["Training logger<br/>[logger.py]"]
end

node_user(("Image and mask user"))
node_output["Inpainted image"]

node_user -->|"provides images/masks"| node_testdata
node_testdata -->|"supplies batches"| node_trainer
node_traindata -->|"supplies batches"| node_trainer
node_trainer -->|"uses line prior"| node_line_detector
node_trainer -->|"uses structure prior"| node_structure_net
node_trainer -->|"uses priors"| node_prior_net
node_trainer -->|"runs generator"| node_generator
node_trainer -->|"runs generator"| node_generator_family
node_trainer -->|"scores images"| node_discriminator
node_trainer -->|"computes losses"| node_losses
node_trainer -->|"produces image"| node_output
node_trainer -->|"logs progress"| node_logger
node_output -.->|"evaluated by"| node_metric
node_trainer -.->|"optional prior"| node_segmentation

click node_testdata "https://github.com/ewrfcas/zits-plusplus/blob/main/dataset/dataloader.py"
click node_traindata "https://github.com/ewrfcas/zits-plusplus/blob/main/base/dataset.py"
click node_trainer "https://github.com/ewrfcas/zits-plusplus/blob/main/trainers/pl_trainers.py"
click node_line_detector "https://github.com/ewrfcas/zits-plusplus/blob/main/trainers/lsm_hawp/detector.py"
click node_segmentation "https://github.com/ewrfcas/zits-plusplus/blob/main/networks/ade20k/base.py"
click node_structure_net "https://github.com/ewrfcas/zits-plusplus/blob/main/networks/tsr.py"
click node_prior_net "https://github.com/ewrfcas/zits-plusplus/blob/main/networks/pcp.py"
click node_generator "https://github.com/ewrfcas/zits-plusplus/blob/main/networks/mat.py"
click node_generator_family "https://github.com/ewrfcas/zits-plusplus/blob/main/networks/generators.py"
click node_discriminator "https://github.com/ewrfcas/zits-plusplus/blob/main/networks/discriminators.py"
click node_losses "https://github.com/ewrfcas/zits-plusplus/blob/main/networks/losses.py"
click node_metric "https://github.com/ewrfcas/zits-plusplus/blob/main/inpainting_metric.py"
click node_logger "https://github.com/ewrfcas/zits-plusplus/blob/main/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_testdata,node_traindata,node_user toneBlue
class node_line_detector,node_segmentation,node_structure_net,node_prior_net toneAmber
class node_generator,node_generator_family,node_discriminator,node_losses toneMint
class node_trainer,node_metric,node_logger toneRose
class node_output toneIndigo
```

## Components

### Input preparation

- [Test dataset](https://github.com/ewrfcas/zits-plusplus/blob/main/dataset/dataloader.py): `dataset/dataloader.py`
- [Training datasets](https://github.com/ewrfcas/zits-plusplus/blob/main/base/dataset.py): `base/dataset.py`

### Structural priors

- [LSM-HAWP detector](https://github.com/ewrfcas/zits-plusplus/blob/main/trainers/lsm_hawp/detector.py): `trainers/lsm_hawp/detector.py`
- [ADE20K segmentation](https://github.com/ewrfcas/zits-plusplus/blob/main/networks/ade20k/base.py): `networks/ade20k/base.py`
- [Structure upsampler](https://github.com/ewrfcas/zits-plusplus/blob/main/networks/tsr.py): `networks/tsr.py`
- [Prior networks](https://github.com/ewrfcas/zits-plusplus/blob/main/networks/pcp.py): `networks/pcp.py`

### Inpainting models

- [MAT generator](https://github.com/ewrfcas/zits-plusplus/blob/main/networks/mat.py): `networks/mat.py`
- [Generator components](https://github.com/ewrfcas/zits-plusplus/blob/main/networks/generators.py): `networks/generators.py`
- [Discriminator](https://github.com/ewrfcas/zits-plusplus/blob/main/networks/discriminators.py): `networks/discriminators.py`
- [Training losses](https://github.com/ewrfcas/zits-plusplus/blob/main/networks/losses.py): `networks/losses.py`

### Runtime and training

- [Inpainting trainer](https://github.com/ewrfcas/zits-plusplus/blob/main/trainers/pl_trainers.py): `trainers/pl_trainers.py`
- [Inpainting metrics](https://github.com/ewrfcas/zits-plusplus/blob/main/inpainting_metric.py): `inpainting_metric.py`
- [Training logger](https://github.com/ewrfcas/zits-plusplus/blob/main/logger/logger.py): `logger/logger.py`

### Other components

- **Image and mask user**
- **Inpainted image**

## Connections

- Image and mask user → Test dataset: provides images/masks
- Test dataset → Inpainting trainer: supplies batches
- Training datasets → Inpainting trainer: supplies batches
- Inpainting trainer → LSM-HAWP detector: uses line prior
- Inpainting trainer → Structure upsampler: uses structure prior
- Inpainting trainer → Prior networks: uses priors
- Inpainting trainer → MAT generator: runs generator
- Inpainting trainer → Generator components: runs generator
- Inpainting trainer → Discriminator: scores images
- Inpainting trainer → Training losses: computes losses
- Inpainting trainer → Inpainted image: produces image
- Inpainting trainer → Training logger: logs progress
- Inpainted image → Inpainting metrics: evaluated by
- Inpainting trainer → ADE20K segmentation: optional prior