malaika16/bag-classifier-v2

Bag Classifier — Classification Head

This repository contains the classification head weights for the Bag Image Classifier, fine-tuned on luxury fashion e-commerce images.

What is this?

This is the final classification layer (head) of a Vision Transformer (ViT) model fine-tuned to classify luxury bag images as accept or reject for e-commerce catalog use.

Base Model

• Architecture: google/vit-base-patch16-224
• Fine-tuning technique: Linear Probing (only head was trained)
• Task: Binary classification (Accept / Reject)

Labels

Label	ID	Meaning
reject	0	Image not suitable for catalog
accept	1	Image suitable for catalog

Accept Categories

Images labeled as accept include:

• Front-idle
• side-idle
• Angle-idle
• outside-shoulder
• outside-hand
• outside-idle

Reject Categories

Images labeled as reject include:

• inside
• Zoom-inside
• Zoom-outside
• Bottom-idle
• OOD (out of distribution — non-bag items)
• etc

Training Data

• 3 datasets combined from different luxury fashion websites
• ~2,232 balanced images (1,116 accept + 1,116 reject)
• 80% train / 20% test split
• Final test accuracy: 98%

How to Use

Step 1 — Install dependencies

pip install transformers torch

Step 2 — Load base model and head

import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification

# Load base ViT model
model = AutoModelForImageClassification.from_pretrained(
    "google/vit-base-patch16-224",
    num_labels=2,
    id2label={0: "reject", 1: "accept"},
    label2id={"reject": 0, "accept": 1},
    ignore_mismatched_sizes=True
)

# Load and apply head weights
head_weights = torch.load("head.pt", map_location="cpu")
model.classifier.weight = head_weights["classifier.weight"]
model.classifier.bias = head_weights["classifier.bias"]

model.eval()
print("Model ready!")

Step 3 — Run inference

from PIL import Image
from transformers import AutoImageProcessor

processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")

# Load your image
image = Image.open("your_bag_image.jpg").convert("RGB")

# Preprocess
inputs = processor(images=image, return_tensors="pt")

# Predict
import torch
with torch.no_grad():
    outputs = model(**inputs)
    predicted_id = outputs.logits.argmax(-1).item()
    label = {0: "reject", 1: "accept"}[predicted_id]
    confidence = torch.softmax(outputs.logits, dim=-1).max().item()

print(f"Prediction: {label}")
print(f"Confidence: {confidence:.2%}")

Training Details

Parameter	Value
Base model	google/vit-base-patch16-224
Epochs	5
Batch size	16
Image size	600x600 (thumbnail)
Optimizer	AdamW (default)
Test accuracy	98%

Dataset

Private dataset hosted on HuggingFace — malaika16/bag-labels

Contains labeled images from 3 luxury fashion websites with 10 categories collapsed into binary accept/reject labels.