|
|
| # 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 |
| ```python |
| pip install transformers torch |
| ``` |
|
|
| ### Step 2 β Load base model and head |
| ```python |
| 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 |
| ```python |
| 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. |
|
|