CoolFace
Modelpublic

0-ma/vit-geometric-shapes-base

sourceHugging Faceapache-2.0updated 1y agoView on Hugging Face
1likes25downloads
Model Card

Model Card for VIT Geometric Shapes Dataset Base

Training Dataset

  • Repository: https://huggingface.co/datasets/0-ma/geometric-shapes

Base Model

  • Repository: https://huggingface.co/google/vit-base-patch16-224

Accuracy

  • Accuracy on dataset 0-ma/geometric-shapes [test] : 0.9993

Loading and using the model

import numpy as np from PIL import Image from transformers import AutoImageProcessor, AutoModelForImageClassification import requests labels = [ "None", "Circle", "Triangle", "Square", "Pentagon", "Hexagon" ] images = [Image.open(requests.get("https://raw.githubusercontent.com/0-ma/geometric-shape-detector/main/input/exemplecircle.jpg", stream=True).raw), Image.open(requests.get("https://raw.githubusercontent.com/0-ma/geometric-shape-detector/main/input/exemplepentagone.jpg", stream=True).raw)] featureextractor = AutoImageProcessor.frompretrained('0-ma/vit-geometric-shapes-base') model = AutoModelForImageClassification.frompretrained('0-ma/vit-geometric-shapes-base') inputs = featureextractor(images=images, returntensors="pt") logits = model(**inputs)['logits'].cpu().detach().numpy() predictions = np.argmax(logits, axis=1) predictedlabels = [labels[prediction] for prediction in predictions] print(predicted_labels)

Model generation

The model has been created using the 'trainshapedetector.py' of the project from the project https://github.com/0-ma/geometric-shape-detector. No external code sources were used.