A11Sunday/vit-cat-dog-panda

ViT cat / dog / panda

Fine-tuned google/vit-base-patch16-224 для классификации изображений на три класса: cat, dog, panda.

Это веса для проекта vit-image-classifier. Модель = кастомный patch embedding + предобученный энкодер ViT + голова-классификатор на [CLS] токене. Стратегия дообучения: linear_probe.

• Validation accuracy: 0.9978
• Классы (порядок индексов): cat, dog, panda
• Вход: изображение RGB, нормализация mean=std=0.5, размер 224×224 → тензор (B, 3, 224, 224).

Файлы

• linear_probe_best.pt — чекпойнт PyTorch (torch.save) со словарём: model_state, class_names, model_name, strategy, epoch, val_acc.

Использование

from huggingface_hub import hf_hub_download
from src.config import Config
from src.inference import Predictor

path = hf_hub_download(repo_id="A11Sunday/vit-cat-dog-panda", filename="linear_probe_best.pt")
predictor = Predictor(Config(), path)
print(predictor.predict(image))  # {"cat": ..., "dog": ..., "panda": ...}

Запуск демо целиком (Gradio) — см. README проекта.