Spaces:
Sleeping
Sleeping
| 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) |