Spaces:
Sleeping
Sleeping
File size: 1,005 Bytes
183c23b c5df48e 183c23b 8ae9b18 183c23b 2d13268 183c23b 2d13268 183c23b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
from PIL import Image
from torchvision import transforms
from transformers import AutoModelForImageClassification
import gradio as gr
import torch
num_classes = 2
def predict(inp):
inputs = data_transforms(inp)[None]
model.eval()
with torch.no_grad():
logits = model(inputs)['logits']
probs = torch.softmax(logits,dim=1)
confidences = {labels[i]: probs[0][i] for i in range(num_classes)}
return confidences
data_transforms = transforms.Compose([
transforms.Resize((224,224)), # Resize the images to a specific size
transforms.ToTensor(), # Convert images to tensors
#transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # Normalize the image data
])
# Load model directly
model = AutoModelForImageClassification.from_pretrained("Manu8/vit_cats-vs-dogs", trust_remote_code=True)
labels = [
'cat','dog'
]
gr.Interface(fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=3)).launch() |