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(...).

piotreksl/vehicle-color-recognition

Vehicle Color Recognition — EfficientNet-B4 (PyTorch)

Model Details

Classes

Usage

Notes

piotreksl/
vehicle-color-recognition