Vehicle Color Recognition — EfficientNet-B4 (PyTorch)

Classifies the dominant colour of a vehicle from a cropped car image.

Model Details

Property	Value
Architecture	efficientnet_b4
Input size	380×380
Classes	14
Format	PyTorch checkpoint (.pth)

Classes

beige, black, blue, brown, gold, green, grey, orange, pink, purple, red, silver, white, yellow

Usage

import torch, timm, yaml
from PIL import Image
import albumentations as A
from albumentations.pytorch import ToTensorV2
import numpy as np

# Load config
with open("config.yaml") as f:
    cfg = yaml.safe_load(f)

# Build model
model = timm.create_model(
    cfg["model_name"], pretrained=False,
    num_classes=cfg["num_classes"], drop_rate=cfg["drop_rate"]
)
ckpt = torch.load("best_model.pth", map_location="cpu", weights_only=False)
model.load_state_dict(ckpt["model_state_dict"])
model.eval()

# Preprocess
tf = A.Compose([
    A.Resize(418, 418, interpolation=2),
    A.CenterCrop(380, 380),
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ToTensorV2(),
])
img = np.array(Image.open("car.jpg").convert("RGB"))
tensor = tf(image=img)["image"].unsqueeze(0)  # [1, 3, H, W]

# Predict
with torch.no_grad():
    probs = torch.softmax(model(tensor), dim=1)[0]
top = probs.argsort(descending=True)[:3]
classes = ['beige', 'black', 'blue', 'brown', 'gold', 'green', 'grey', 'orange', 'pink', 'purple', 'red', 'silver', 'white', 'yellow']
for i in top:
    print(f"{classes[i]}: {probs[i]:.2%}")

Notes

Input images should be cropped to the car area before prediction.
tan is merged into beige during training; there is no separate tan class.
For CPU-only / production use, export to ONNX from this checkpoint in your downstream consumer: torch.onnx.export(...).

Downloads last month: -