from transformers import AutoImageProcessor, AutoModelForImageClassification
import numpy as np
import torch
import gradio as gr


model = AutoModelForImageClassification.from_pretrained('hero_photo_eligibility_model')
checkpoint = 'google/vit-base-patch16-224'
image_processor = AutoImageProcessor.from_pretrained(checkpoint)

label_names = ['NO', 'YES']

def classify(im):
    features = image_processor(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(label_names)}
    return confidences


title = """Detecting whether a photo is suitable for VDP main photo"""

description = """Hero photo eligibility demo"""

interface = gr.Interface(
      fn=classify,
      inputs='image',
      outputs='label',
      title=title,
      description=description
  )

interface.launch(share=True, debug=True)