from transformers import ViTImageProcessor, ViTForImageClassification
import torch
import gradio as gr

feature_extractor = ViTImageProcessor.from_pretrained("model_artifacts")
model = ViTForImageClassification.from_pretrained("model_artifacts")

labels = ['Chevrolet Equinox',
 'Chevrolet Silverado 1500',
 'Ford Escape',
 'Ford Explorer',
 'Ford F-150',
 'GMC Sierra 1500',
 'Honda CR-V',
 'Jeep Compass',
 'Jeep Grand Cherokee',
 'Jeep Wrangler',
 'Mazda CX-5',
 'Nissan Rogue',
 'RAM 1500',
 'RAM 2500',
 'Toyota Camry']

def classify(im):
  features = feature_extractor(im, return_tensors='pt')
  logits = model(features["pixel_values"])[-1]
  probability = torch.nn.functional.softmax(logits, dim=-1)
  probs = probability[0].detach().numpy()
  confidences = {label: float(probs[i]) for i, label in enumerate(labels)} 
  return confidences


description = """
        Simple car recognition model. Can recognize one of the followings: 

        Chevrolet Equinox
        Chevrolet Silverado 1500
        Ford Escape
        Ford Explorer
        Ford F-150
        GMC Sierra 1500
        Honda CR-V
        Jeep Compass
        Jeep Grand Cherokee
        Jeep Wrangler
        Mazda CX-5
        Nissan Rogue
        RAM 1500
        RAM 2500
        Toyota Camry
        
        """
interface = gr.Interface(fn=classify,
                         inputs="image",
                         outputs="label",
                         title="Car classification demo :)",
                         description=description ) 

interface.launch()