# arshad1428/image-classification-cnn architecture

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

- Interactive diagram: https://gitdiagram.com/arshad1428/image-classification-cnn
- Repository: https://github.com/arshad1428/image-classification-cnn

## Overview

This project trains and serves six-class scene image classifiers. The main workflow loads and splits Intel images, applies training/evaluation transforms, trains a selected CNN, validates and checkpoints it, then evaluates on the held-out test set while recording metrics in MLflow. A separate Streamlit workflow loads a checkpoint and predicts a class and confidence for an uploaded image. The README documents both model variants and experiment tracking; unsampled implementation details are represented only where the documented capability and tree support them.

Main entry points: [`main.py`](https://github.com/arshad1428/image-classification-cnn/blob/main/main.py) (Training entry), [`app.py`](https://github.com/arshad1428/image-classification-cnn/blob/main/app.py) (Streamlit app), [`src/image_classifier/models/transfer_learning.py`](https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/models/transfer_learning.py) (ResNet18 model), [`src/image_classifier/inference/predict.py`](https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/inference/predict.py) (Image prediction)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Image data"]
  node_raw_images["Intel images"]
  node_loader["Dataset splits<br/>[data_loader.py]"]
  node_transforms["Image transforms<br/>[preprocessing.py]"]
end

subgraph group_training["Training workflow"]
  node_entry["Training entry<br/>[main.py]"]
  node_train_loops["Train and validate<br/>[train.py]"]
  node_config["Project config<br/>[config.py]"]
end

subgraph group_models["Models and artifacts"]
  node_cnn["Custom CNN<br/>[cnn.py]"]
  node_resnet["ResNet18 model"]
  node_checkpoint["Checkpoint handling<br/>[checkpoint.py]"]
  node_model_files[("Saved checkpoints")]
end

subgraph group_serving["Evaluation and serving"]
  node_evaluation["Test metrics<br/>[evaluate.py]"]
  node_tracking[("MLflow tracking")]
  node_app["Streamlit app<br/>[app.py]"]
  node_predict["Image prediction<br/>[predict.py]"]
end

node_user(("User"))
node_pretrained[("Torchvision weights")]

node_user -->|"uploads image"| node_app
node_raw_images -->|"reads images"| node_loader
node_loader -->|"gets transforms"| node_transforms
node_loader -->|"returns loaders"| node_entry
node_config -->|"supplies paths"| node_loader
node_config -->|"supplies settings"| node_entry
node_entry -->|"creates loaders"| node_loader
node_entry -->|"constructs when selected"| node_cnn
node_entry -->|"constructs when selected"| node_resnet
node_resnet -.->|"loads weights"| node_pretrained
node_entry -->|"runs epochs"| node_train_loops
node_train_loops -->|"calls model"| node_cnn
node_train_loops -->|"calls model"| node_resnet
node_entry -->|"saves best model"| node_checkpoint
node_checkpoint -->|"writes checkpoint"| node_model_files
node_entry -->|"evaluates test set"| node_evaluation
node_evaluation -->|"calls model"| node_cnn
node_evaluation -->|"calls model"| node_resnet
node_entry -->|"logs params and metrics"| node_tracking
node_entry -->|"logs checkpoint artifact"| node_tracking
node_app -->|"constructs when selected"| node_cnn
node_app -->|"constructs when selected"| node_resnet
node_app -->|"loads checkpoint"| node_model_files
node_app -->|"requests prediction"| node_predict
node_predict -->|"gets eval transform"| node_transforms
node_predict -->|"calls model"| node_cnn
node_predict -->|"calls model"| node_resnet
node_predict -->|"returns class and confidence"| node_app
node_app -->|"shows result"| node_user

click node_raw_images "https://github.com/arshad1428/image-classification-cnn/tree/main/data/raw"
click node_loader "https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/data_loader.py"
click node_transforms "https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/preprocessing.py"
click node_entry "https://github.com/arshad1428/image-classification-cnn/blob/main/main.py"
click node_train_loops "https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/training/train.py"
click node_config "https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/config.py"
click node_cnn "https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/models/cnn.py"
click node_resnet "https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/models/transfer_learning.py"
click node_checkpoint "https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/training/checkpoint.py"
click node_model_files "https://github.com/arshad1428/image-classification-cnn/tree/main/models"
click node_evaluation "https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/evaluation/evaluate.py"
click node_app "https://github.com/arshad1428/image-classification-cnn/blob/main/app.py"
click node_predict "https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/inference/predict.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_raw_images,node_loader,node_transforms,node_user toneBlue
class node_entry,node_train_loops,node_config,node_pretrained toneAmber
class node_cnn,node_resnet,node_checkpoint,node_model_files toneMint
class node_evaluation,node_tracking,node_app,node_predict toneRose
```

## Components

### Image data

- [Intel images](https://github.com/arshad1428/image-classification-cnn/tree/main/data/raw): `data/raw`
- [Dataset splits](https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/data_loader.py): `src/image_classifier/data_loader.py`
- [Image transforms](https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/preprocessing.py): `src/image_classifier/preprocessing.py`

### Training workflow

- [Training entry](https://github.com/arshad1428/image-classification-cnn/blob/main/main.py): `main.py`
- [Train and validate](https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/training/train.py): `src/image_classifier/training/train.py`
- [Project config](https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/config.py): `src/image_classifier/config.py`

### Models and artifacts

- [Custom CNN](https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/models/cnn.py): `src/image_classifier/models/cnn.py`
- [ResNet18 model](https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/models/transfer_learning.py): `src/image_classifier/models/transfer_learning.py`
- [Checkpoint handling](https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/training/checkpoint.py): `src/image_classifier/training/checkpoint.py`
- [Saved checkpoints](https://github.com/arshad1428/image-classification-cnn/tree/main/models): `models`

### Evaluation and serving

- [Test metrics](https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/evaluation/evaluate.py): `src/image_classifier/evaluation/evaluate.py`
- **MLflow tracking**
- [Streamlit app](https://github.com/arshad1428/image-classification-cnn/blob/main/app.py): `app.py`
- [Image prediction](https://github.com/arshad1428/image-classification-cnn/blob/main/src/image_classifier/inference/predict.py): `src/image_classifier/inference/predict.py`

### Other components

- **User**
- **Torchvision weights**

## Connections

- User → Streamlit app: uploads image
- Intel images → Dataset splits: reads images
- Dataset splits → Image transforms: gets transforms
- Dataset splits → Training entry: returns loaders
- Project config → Dataset splits: supplies paths
- Project config → Training entry: supplies settings
- Training entry → Dataset splits: creates loaders
- Training entry → Custom CNN: constructs when selected
- Training entry → ResNet18 model: constructs when selected
- ResNet18 model → Torchvision weights: loads weights
- Training entry → Train and validate: runs epochs
- Train and validate → Custom CNN: calls model
- Train and validate → ResNet18 model: calls model
- Training entry → Checkpoint handling: saves best model
- Checkpoint handling → Saved checkpoints: writes checkpoint
- Training entry → Test metrics: evaluates test set
- Test metrics → Custom CNN: calls model
- Test metrics → ResNet18 model: calls model
- Training entry → MLflow tracking: logs params and metrics
- Training entry → MLflow tracking: logs checkpoint artifact
- Streamlit app → Custom CNN: constructs when selected
- Streamlit app → ResNet18 model: constructs when selected
- Streamlit app → Saved checkpoints: loads checkpoint
- Streamlit app → Image prediction: requests prediction
- Image prediction → Image transforms: gets eval transform
- Image prediction → Custom CNN: calls model
- Image prediction → ResNet18 model: calls model
- Image prediction → Streamlit app: returns class and confidence
- Streamlit app → User: shows result