hair_color_classifier

Hair Color ConvNeXt Tiny

This is a seven-class hair-color image classifier fine-tuned from timm/convnext_tiny.fb_in22k. It predicts one of the following labels, in model-output order:

Model details

The classifier head is a dropout layer with probability 0.15 followed by a seven-output linear layer. The backbone and classifier were fine-tuned for the hair-color task. The published artifacts are full-precision exports, not quantized variants. The ONNX model supports dynamic batch sizes.

Files

• hair_color-convnext_tiny.fb_in22k.onnx is the portable inference graph.
• hair_color-convnext_tiny.fb_in22k.safetensors contains the PyTorch HairClassifier state dictionary.
• config.json records the architecture, labels, and timm model settings.
• preprocessor_config.json records the image preprocessing contract.

The safetensors keys include the custom wrapper's model. prefix and custom classifier head. Load them with the HairClassifier implementation from the source repository, not as an unmodified upstream timm checkpoint.

Preprocessing

For the model tensor itself:

1. Convert the image to RGB.
2. Resize directly to 224 × 224 using bilinear interpolation.
3. Rescale unsigned 8-bit pixels by 1 / 255.
4. Normalize channels with ImageNet mean [0.485, 0.456, 0.406] and standard deviation [0.229, 0.224, 0.225].
5. Convert from HWC to NCHW layout and add the batch dimension.

The source application first detects and aligns a face with InsightFace buffalo_l, retaining 20% padding around the aligned face. Inputs that are already face-centered or similarly aligned are closest to the training and application pipeline.

ONNX inference

Install huggingface_hub, onnxruntime, numpy, and Pillow, then run:

from huggingface_hub import hf_hub_download
import numpy as np
import onnxruntime as ort
from PIL import Image

REPO_ID = "electblake/hair_color_classifier"
LABELS = ["black", "blonde", "blue", "brown", "pink", "red", "silver"]
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)

model_path = hf_hub_download(
    REPO_ID,
    "hair_color-convnext_tiny.fb_in22k.onnx",
)
image = Image.open("hair.jpg").convert("RGB")
image = image.resize((224, 224), Image.Resampling.BILINEAR)
image = np.asarray(image, dtype=np.float32) / 255.0
image = (image - MEAN) / STD
image = np.transpose(image, (2, 0, 1))[None, ...]

session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
logits = session.run(["logits"], {"image": image})[0][0]
probabilities = np.exp(logits - logits.max())
probabilities /= probabilities.sum()

prediction = int(probabilities.argmax())
print(LABELS[prediction], float(probabilities[prediction]))

Safetensors inference

From a checkout of the source repository:

from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

from src.model import HairClassifier

weights_path = hf_hub_download(
    "electblake/hair_color_classifier",
    "hair_color-convnext_tiny.fb_in22k.safetensors",
)
model = HairClassifier(
    model_name="convnext_tiny.fb_in22k",
    num_classes=7,
    dropout=0.15,
    pretrained=False,
)
model.load_state_dict(load_file(weights_path))
model.eval()

Training data

The training corpus contains 29,662 images across the seven labels. It combines CelebAMask-HQ images selected by mutually exclusive hair-color attributes with additional locally supplied class folders. The additional sources are aligned before splitting; the source datasets themselves are not redistributed here.

The split ratio is 70%/15%/15% with random seed 42. Training uses weighted cross-entropy, AdamW, MixUp, CutMix, horizontal flips, geometric transforms, brightness/contrast changes, hue/saturation changes, and Gaussian noise. The backbone is frozen for the first 12 epochs.

Evaluation

The published checkpoint was selected at epoch 28 with 94.88% validation accuracy. This is a model-selection result on the project's validation split, not an independent benchmark or a test-set result.

Performance is expected to vary with lighting, color casts, occlusion, wigs, dyed or multicolored hair, grayscale images, unusual crops, and failed face alignment. Confidence scores are softmax probabilities and have not been calibrated.

Intended use and limitations

The model is intended for research, media organization, and non-critical hair-color tagging. It is not an identity model and should not be used for biometrics, surveillance, demographic inference, or decisions affecting a person's rights or access to services.

Training data may not represent all skin tones, ages, hairstyles, cultural contexts, cameras, or lighting conditions evenly. Evaluate the model on the target population and setting before use.