# santiagosabogall/gan-for-galaxy-image-generation architecture

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

- Interactive diagram: https://gitdiagram.com/santiagosabogall/gan-for-galaxy-image-generation
- Repository: https://github.com/santiagosabogall/gan-for-galaxy-image-generation

## Overview

1. This repository is a single-purpose TensorFlow project for training a **Wasserstein GAN with Gradient Penalty (WGAN-GP)** to generate **64×64 RGB galaxy images** from the Galaxy Zoo 2 dataset. It is an offline model-training pipeline, not a serving or application backend.

2. The code is split into five clear subsystems under `src/`:  
   - `config.py` for global hyperparameters and paths  
   - `data.py` for dataset acquisition and preprocessing  
   - `models.py` for Generator and Critic definitions  
   - `train.py` for the custom WGAN-GP training loop  
   - `utils.py` for visualization, GIF creation, and plotting

3. The main execution entrypoint is `python -m src.train`. That script orchestrates the full lifecycle: load configuration, obtain data, instantiate models, run alternating Critic/Generator updates, and emit artifacts during training.

4. The **data boundary** is external: the project downloads the Galaxy Zoo 2 dataset via **kagglehub** and feeds it through a **tf.data** pipeline. This means the training loop consumes an abstracted TensorFlow dataset rather than reading raw files directly.

5. The **preprocessing boundary** likely normalizes images into the range expected by the Generator/Critic pair. The README indicates a `tanh` output on the Generator, so the image pipeline almost certainly maps real images into a matching normalized scale for WGAN training.

6. The **model subsystem** contains two networks with distinct roles:  
   - **Generator**: latent vector `z ∈ R^128` → synthetic `64×64×3` image  
   - **Critic**: `64×64×3` image → scalar Wasserstein score  
   The architecture uses transposed convolutions in the Generator and strided convolutions in the Critic.

7. The WGAN-GP-specific architectural constraint is important: the Critic intentionally avoids BatchNormalization, because gradient penalty depends on well-behaved sample-wise gradients. That choice is a key boundary between “ordinary CNN classifier” design and WGAN critic design.

8. The **training subsystem** is a custom adversarial optimization loop rather than a standard Keras `fit()` workflow. It alternates Critic and Generator updates and computes the gradient penalty term explicitly, which is central to the repository’s behavior.

9. The repository also has a clear **artifact/output boundary**. During training, it writes:  
   - model checkpoints into `checkpoints/`  
   - preview images into `output/`  
   - TensorBoard logs into `logs/`  
   These outputs are the main integration points for monitoring and resuming experiments.

10. `utils.py` provides the post-training/monitoring layer: it turns epoch snapshots into visual progress assets, including an animated GIF that summarizes training evolution. This is a presentation/inspection concern, separate from the core optimization logic.

11. `requirements.txt` and the README indicate the runtime stack is **Python 3.9+** with **TensorFlow 2.15+**. External tooling that materially affects the architecture includes **TensorBoard** for metrics inspection and **kagglehub** for dataset retrieval/caching.

12. The notebook `GAN_images_of_galaxies.ipynb` is a legacy or exploratory artifact, while `src/` is the productionized training implementation. In an architecture diagram, the notebook should be shown as optional historical context, not part of the main runtime path.

13. In diagram terms, the primary flow is:  
   **Kaggle/Galaxy Zoo 2 → data pipeline → training loop → Generator/Critic updates → checkpoints/logs/images/GIFs**.  
   The only persistent external dependencies are the dataset source and the local artifact directories used for experiment tracking.

Main entry points: [`src/train.py`](https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/train.py) (WGAN-GP loop), [`src/data.py`](https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/data.py) (Data pipeline), [`src/models.py`](https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/models.py) (Models), [`src/utils.py`](https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/utils.py) (Utils)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_source["Source code"]
  node_train_entry["Train entry<br/>[train.py]"]
  node_config["Config<br/>settings<br/>[config.py]"]
  node_data["Data pipeline<br/>input pipeline<br/>[data.py]"]
  node_models["Models<br/>networks<br/>[models.py]"]
  node_utils["Utils<br/>visualization<br/>[utils.py]"]
  node_generator{{"Generator<br/>network<br/>[models.py]"}}
  node_critic{{"Critic<br/>network<br/>[models.py]"}}
  node_wgan_loop["WGAN-GP loop<br/>training loop<br/>[train.py]"]
  node_gradient_penalty["Gradient penalty<br/>loss term<br/>[train.py]"]
  node_notebook["Notebook"]
end

subgraph group_external["External dependencies"]
  node_kagglehub[("KaggleHub<br/>dataset fetch")]
  node_tfdata["tf.data<br/>pipeline api"]
  node_galaxy_zoo2[("Galaxy Zoo 2<br/>dataset")]
  node_tensorboard["TensorBoard<br/>metrics viewer"]
end

subgraph group_artifacts["Training artifacts"]
  node_checkpoints[("Checkpoints<br/>artifact dir")]
  node_output[("Output images<br/>artifact dir")]
  node_logs[("TensorBoard logs<br/>artifact dir")]
  node_gif["Training GIF<br/>artifact"]
end

node_train_entry -->|"loads"| node_config
node_train_entry -->|"prepares"| node_data
node_data -->|"downloads"| node_kagglehub
node_kagglehub -->|"sources"| node_galaxy_zoo2
node_data -->|"yields"| node_tfdata
node_train_entry -->|"instantiates"| node_models
node_models -->|"defines"| node_generator
node_models -->|"defines"| node_critic
node_train_entry -->|"runs"| node_wgan_loop
node_wgan_loop -->|"consumes"| node_tfdata
node_wgan_loop -->|"updates"| node_critic
node_wgan_loop -->|"updates"| node_generator
node_wgan_loop -->|"computes"| node_gradient_penalty
node_wgan_loop -->|"writes"| node_checkpoints
node_wgan_loop -->|"writes"| node_logs
node_wgan_loop -->|"writes"| node_output
node_logs -->|"viewed in"| node_tensorboard
node_utils -->|"renders"| node_output
node_utils -->|"assembles"| node_gif
node_train_entry -->|"uses"| node_utils
node_notebook -.->|"prototypes"| node_models
node_notebook -.->|"prototypes"| node_data

click node_train_entry "https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/train.py"
click node_config "https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/config.py"
click node_data "https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/data.py"
click node_models "https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/models.py"
click node_utils "https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/utils.py"
click node_generator "https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/models.py"
click node_critic "https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/models.py"
click node_wgan_loop "https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/train.py"
click node_gradient_penalty "https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/train.py"
click node_notebook "https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/GAN_images_of_galaxies.ipynb"

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_entry,node_config,node_data,node_models,node_utils,node_generator,node_critic,node_wgan_loop,node_gradient_penalty,node_notebook toneBlue
class node_kagglehub,node_tfdata,node_galaxy_zoo2,node_tensorboard toneAmber
class node_checkpoints,node_output,node_logs,node_gif toneMint
```

## Components

### Source code

Core TensorFlow training implementation under src/

- [Train entry](https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/train.py): module · Main orchestration entrypoint for WGAN-GP training · `src/train.py`
- [Config](https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/config.py): settings · Global hyperparameters, paths, and run settings · `src/config.py`
- [Data pipeline](https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/data.py): input pipeline · Dataset acquisition, preprocessing, and tf.data construction · `src/data.py`
- [Models](https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/models.py): networks · Generator and critic definitions · `src/models.py`
- [Utils](https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/utils.py): visualization · Plotting, previews, and GIF creation · `src/utils.py`
- [Generator](https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/models.py): network · Maps latent vectors to 64x64 RGB galaxy images · `src/models.py`
- [Critic](https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/models.py): network · Scores real and generated images for Wasserstein training · `src/models.py`
- [WGAN-GP loop](https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/train.py): training loop · Alternating critic and generator updates with custom training logic · `src/train.py`
- [Gradient penalty](https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/src/train.py): loss term · Explicit WGAN-GP regularization computed in the loop · `src/train.py`
- [Notebook](https://github.com/santiagosabogall/gan-for-galaxy-image-generation/blob/main/GAN_images_of_galaxies.ipynb): legacy notebook · Exploratory or historical implementation artifact · `GAN_images_of_galaxies.ipynb`

### External dependencies

Dataset retrieval and experiment tooling outside the repo

- **KaggleHub**: dataset fetch · Downloads and caches the Galaxy Zoo 2 dataset
- **tf.data**: pipeline api · TensorFlow dataset abstraction consumed by training
- **Galaxy Zoo 2**: dataset · Source dataset of galaxy images
- **TensorBoard**: metrics viewer · Inspects training summaries and logged scalars

### Training artifacts

Files and directories produced during training

- **Checkpoints**: artifact dir · Saved model state for resuming training
- **Output images**: artifact dir · Preview images generated during training
- **TensorBoard logs**: artifact dir · Scalar and summary data for monitoring
- **Training GIF**: artifact · Animated summary of training progress

## Connections

- Train entry → Config: loads · Reads hyperparameters, paths, and training settings
- Train entry → Data pipeline: prepares · Builds the dataset input pipeline
- Data pipeline → KaggleHub: downloads · Fetches Galaxy Zoo 2 through KaggleHub
- KaggleHub → Galaxy Zoo 2: sources · External dataset source
- Data pipeline → tf.data: yields · Produces TensorFlow datasets for training
- Train entry → Models: instantiates · Creates generator and critic networks
- Models → Generator: defines · Generator architecture and output scaling
- Models → Critic: defines · Critic architecture without batch normalization
- Train entry → WGAN-GP loop: runs · Executes the custom adversarial optimization loop
- WGAN-GP loop → tf.data: consumes · Reads batches from the dataset pipeline
- WGAN-GP loop → Critic: updates · Trains critic on real and generated images
- WGAN-GP loop → Generator: updates · Trains generator to fool the critic
- WGAN-GP loop → Gradient penalty: computes · Adds WGAN-GP regularization to the critic loss
- WGAN-GP loop → Checkpoints: writes · Persists model checkpoints
- WGAN-GP loop → TensorBoard logs: writes · Emits TensorBoard summaries
- WGAN-GP loop → Output images: writes · Saves preview images during training
- TensorBoard logs → TensorBoard: viewed in · Monitored with TensorBoard
- Utils → Output images: renders · Creates visual samples and progress images
- Utils → Training GIF: assembles · Combines epoch snapshots into an animation
- Train entry → Utils: uses · Calls visualization and artifact helpers
- Notebook → Models: prototypes · Legacy exploratory model code path
- Notebook → Data pipeline: prototypes · Legacy exploratory data handling