Skip to content

README.md

Latest commit

 

History

History
115 lines (94 loc) · 3.93 KB

README.md

File metadata and controls

115 lines (94 loc) · 3.93 KB
Raw

Art Style Classifier API

ML-backed API that predicts an artwork’s style from an input image (27 WikiArt-style categories). It also includes an optional Postgres-backed “gallery” endpoint for serving example artworks by style.

For recruiters

This repo is a compact example of:

  • Serving a PyTorch model behind a production-friendly FastAPI service (lazy model loading, simple input validation, Dockerized deployment).
  • Basic backend hygiene (health endpoint, env-driven config, non-leaky error handling).
  • A small data pipeline (seed_wikiart.py) to populate a Postgres table from a public dataset, storing images in Cloudinary for fast delivery.

Tech stack

  • API: Python 3.11, FastAPI, Uvicorn
  • ML inference: PyTorch, torchvision (EfficientNet-B3), Pillow
  • Data: Postgres (psycopg2)
  • Optional seeding: HuggingFace datasets, tqdm, Cloudinary

API endpoints

Method Path Description
GET /health Liveness check
GET /styles Returns the supported style labels
POST /predict-style Upload an image and get top-k style predictions
GET /gallery Returns artworks matching a style (requires Postgres)

POST /predict-style

  • Body: multipart/form-data with field file (image/jpeg, image/png, or image/webp)
  • Query params: top_k (default 5)
  • Response: the best prediction plus the top-k list with probabilities

Example:

curl -F "file=@/path/to/image.jpg" "http://localhost:10000/predict-style?top_k=5"

GET /gallery (optional)

Returns random artworks for a given style.

  • Query params:
    • style (required) — style/category name (e.g., Impressionism)
    • limit (default 24, min 1, max 60)
    • exclude_id (optional) — omit a specific artwork id
  • Requires: DATABASE_URL

This endpoint expects a table: artworks(id, title, artist, style, image_url)

Running locally

1) Install dependencies

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

2) Provide model weights

At startup the API ensures models/efficientnet_b3_best.pt exists:

  • Recommended: set MODEL_URL to a downloadable .pt file (the server downloads it on first run).
  • Or: place the weights file at models/efficientnet_b3_best.pt.

3) Start the server

uvicorn app:app --reload --host 0.0.0.0 --port 10000

4) Quick checks

curl http://localhost:10000/health
curl http://localhost:10000/styles

Docker

docker build -t art-style-api .
docker run -p 10000:10000 -e MODEL_URL="https://..." art-style-api

Gallery seeding (optional)

seed_wikiart.py can populate the artworks table used by /gallery by:

  1. pulling images from the HuggingFace WikiArt dataset, then
  2. uploading them to Cloudinary, and
  3. upserting metadata into Postgres.

This script uses extra dependencies that aren’t in requirements.txt:

pip install datasets tqdm cloudinary

Required env vars:

  • DATABASE_URL
  • CLOUDINARY_CLOUD_NAME
  • CLOUDINARY_API_KEY
  • CLOUDINARY_API_SECRET

Optional env vars:

  • PER_STYLE (default 300)
  • CLOUDINARY_FOLDER (default wikiart)
  • WIKIART_SPLIT (default train)
  • STYLE_WHITELIST_FILE (optional path to idx_to_style.json, a JSON list, or a newline-separated list of styles)

Run:

python seed_wikiart.py

Project layout

  • app.py: FastAPI routes (/predict-style, /gallery, etc.)
  • model.py: model loading + preprocessing + prediction helpers
  • idx_to_style.json: class label map
  • seed_wikiart.py: dataset → Cloudinary + Postgres seeding pipeline
  • Dockerfile: container build for deployment

Notes / next improvements

  • CORS is currently wide-open for easier demo; production should restrict allow_origins.
  • /gallery uses a new DB connection per request; a connection pool would be the next step for higher throughput.
  • If exposing publicly: add auth, rate limiting, and request size limits.