laszlokiss27/doodle-dash2

Model description

This is a fine-tuned model based on apple/mobilevitv2-1.0-imagenet1k-256 trained for sketch image recognition using Xenova/quickdraw-small dataset.

How to use?

from transformers import MobileViTImageProcessor, MobileViTV2ForImageClassification
from PIL import Image
import requests
import torch
import numpy as np  # Importing NumPy

url = "https://static.thenounproject.com/png/2024184-200.png"
response = requests.get(url, stream=True)

# Convert to grayscale to ensure a single channel input
image = Image.open(response.raw).convert('L')  # Convert to grayscale

processor = MobileViTImageProcessor.from_pretrained("laszlokiss27/doodle-dash2")
model = MobileViTV2ForImageClassification.from_pretrained("laszlokiss27/doodle-dash2")

# Convert the PIL image to a tensor and add a channel dimension
image_tensor = torch.unsqueeze(torch.tensor(np.array(image)), 0).float()
image_tensor = image_tensor.unsqueeze(0)  # Add batch dimension

# Check if processor requires specific form of input
inputs = processor(images=image_tensor, return_tensors="pt")

outputs = model(**inputs)
logits = outputs.logits

# Get prediction
predicted_class_idx = logits.argmax(-1).item()
predicted_class = model.config.id2label[predicted_class_idx]
print("Predicted class:", predicted_class)