0-ma/mit-b0-geometric-shapes

Model Card for Mit-B0 Geometric Shapes Dataset

Training Dataset

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

Base Model

• —Repository: https://huggingface.co/nvidia/mit-b0

Accuracy

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

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/mit-b0-geometric-shapes') model = AutoModelForImageClassification.frompretrained('0-ma/mit-b0-geometric-shapes') 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)

License

The license for this model can be found here.

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.