VIT_10000_2 / app.py
adwod's picture
Update app.py
c861295
import streamlit as st
from transformers import pipeline
from PIL import Image
from transformers import (
AutoImageProcessor,
AutoModelForImageClassification
)
import torch.nn.functional as F
import numpy as np
pre_process = AutoImageProcessor.from_pretrained('adwod/VIT_10000')
st.title("Veo,veo....")
file_name = st.file_uploader("Upload a candidate image")
if file_name is not None:
col1, col2 = st.columns(2)
image = Image.open(file_name)
col1.image(image, use_column_width=True)
inputs = pre_process(images=image, return_tensors="pt")
input_pixels = inputs.pixel_values
model = AutoModelForImageClassification.from_pretrained('adwod/VIT_10000')
outputs = model(input_pixels)
probs = F.softmax(outputs.logits, dim=-1).detach().cpu().numpy()[0]
labels = model.config.id2label.values()
label_probs = zip(labels, probs)
for label, prob in label_probs:
col2.write(f"{label}: {prob:.2f}")
col2.header(model.config.id2label[outputs.logits.argmax(-1).item()])