fanclan/age-gender-model

Image to Age & Gender Model

This repository contains tools for converting, quantizing, and deploying an image-based age and gender prediction model to Hugging Face Inference API.

Model Details

The model is a multi-task ResNet50 architecture that predicts both age group and gender from profile images. It outputs:

• Age classification into 5 age groups (e.g., 18-24, 25-34, 35-44, 45-54, 55+)
• Gender classification (Woman/Man)

Quick Start

Follow these steps to convert your PyTorch model to ONNX, quantize it, and deploy to Hugging Face:

1. Convert PyTorch model to ONNX

python convert_to_onnx.py --pth-path model.pth --onnx-path model.onnx --num-age-classes 5 --num-gender-classes 2

2. Quantize ONNX model

python quantize_model.py --input-path model.onnx --output-path model_quantized.onnx

This reduces model size and improves inference speed.

3. Test locally with FastAPI

python app.py

Visit http://localhost:8000/docs in your browser to test the API.

Deployment Options

Option 1: Hugging Face Inference API

1. Create a new Hugging Face model repository

2. Upload the following files:

   • model_quantized.onnx
   • handler.py
   • requirements.txt
   • README.md
   • .huggingface/inference-config.yaml

3. Set the following environment variable in your Hugging Face model settings:

   • AGE_CLASSES: Comma-separated list of age class names (e.g., "18-24,25-34,35-44,45-54,55+")

Option 2: Hugging Face Spaces (Gradio UI)

This option provides a web UI for your model:

1. Create a new Hugging Face Space

2. Upload the following files:

   • model_quantized.onnx
   • handler.py
   • app_huggingface.py (rename to app.py when uploading)
   • requirements.txt
   • README.md

3. In the Space settings, set:

   • SDK: Gradio
   • Python version: 3.10

API Usage

Input

• Image file (JPG/PNG)

Output

{
  "age": {
    "prediction": "<AGE_GROUP>",
    "confidence": 0.85,
    "all_probabilities": {
      "<AGE_GROUP_1>": 0.05,
      "<AGE_GROUP_2>": 0.85,
      ...
    }
  },
  "gender": {
    "prediction": "<GENDER>",
    "confidence": 0.92,
    "all_probabilities": {
      "Woman": 0.08,
      "Man": 0.92
    }
  }
}

Files Overview

• convert_to_onnx.py: Converts PyTorch model to ONNX format
• model.py: Model architecture definition
• quantize_model.py: Quantizes ONNX model for faster inference
• handler.py: Inference handler for Hugging Face deployment
• app.py: FastAPI application for local testing
• app_huggingface.py: Gradio UI for Hugging Face Spaces
• requirements.txt: Required Python dependencies
• Dockerfile: Container definition for deployment