Update README.md

1ae9e7b verified 11 months ago

1.37 kB

library_name: transformers
license: apache-2.0
base_model: google/vit-base-patch16-224
tags:
  - generated_from_trainer
model-index:
  - name: results
    results: []

E-commerce Classification Model

This model is a fine-tuned version of google/vit-base-patch16-224.

The training and evaluation data consists of product images scraped from e-commerce websites across 50 e-commerce categories. The data was then split into a training set (80%) and a test set (20%). During preprocessing, each image was converted to RGB and resized to 224x224 pixels.

See below to get started:

from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image

model = AutoModelForImageClassification.from_pretrained("khayreali/vit-base-ecommerce-classification")
processor = AutoImageProcessor.from_pretrained("khayreali/vit-base-ecommerce-classification")

clf = pipeline('image-classification',
               model=model,
               image_processor=processor)
     

test_img = Image.open('IMAGE')
     

result = clf(test_img)

# Top predicted category
print(result[0]['label'])