kevincluo's picture
Update app.py
2327eb7
raw
history blame
745 Bytes
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
extractor = AutoFeatureExtractor.from_pretrained("saved_model_files")
model = AutoModelForImageClassification.from_pretrained("saved_model_files")
labels = ['affected', 'destroyed', 'no_damage']
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(labels)}
return confidences
import gradio as gr
interface = gr.Interface(
fn=classify,
inputs="image",
outputs="label"
)
interface.launch(share=True, debug=True)