title	sdk	emoji	colorFrom	colorTo
flowers-classification	docker	🌍	gray	green

About

Image classification model trained using PyTorch Lightning framework and shared on Hugging Face with the use of gradio and Docker.

• Task: Classification of flowers images into one of 102 categories
• Dataset: 102 Category Flower Dataset
• Architecture: SqueezeNet, loaded pretrained ImageNet weights and finetuned on 102 Category Flower Dataset

Tech stack

• PyTorch - neural networks architectures and datasets classes
• PyTorch Lightning - model training and evaluation
• plotly - visualizations
• WandB - metrics, visualizations and model logging
• torchmetrics - metrics calculation
• gradio - application used to show how model works in real world
• Docker - containerize application to allow for Hugging Face deploy

Commands

Three possible ways to go with this repository, all require to start from the repository root path.

Model training

After the model training, two files are saved in "model" directory, transform.pt - torch transform used to transform input image, model.pt - torch model used to predict flower species.

1. Go to backend service

cd services/backend

2. Install dependencies with poetry

poetry install

3. Create virtual environment with poetry

poetry shell

4. Run python script to train the model

python src/train.py

---

Gradio demo app

1. Build docker image and run container with gradio app:

docker build -t flowers .
docker run -it -p 7860:7860 --name flowers_app flowers

---

FastAPI backend with ReactJS frontend

1. Run backend and frontend services with docker-compose

docker-compose up --build

2. Open http://0.0.0.0:5000/docs in your browser to see possible FastAPI endpoints
3. Open http://0.0.0.0:3000 in your browser to see ReactJS frontend