|
|
import streamlit as st |
|
|
from transformers import pipeline |
|
|
from PIL import Image |
|
|
|
|
|
|
|
|
@st.cache_resource |
|
|
def load_image_classifier(): |
|
|
return pipeline("image-classification", model="google/vit-base-patch16-224") |
|
|
|
|
|
@st.cache_resource |
|
|
def load_text_classifier(): |
|
|
return pipeline("sentiment-analysis") |
|
|
|
|
|
|
|
|
image_classifier = load_image_classifier() |
|
|
text_classifier = load_text_classifier() |
|
|
|
|
|
|
|
|
st.title("Hugging Face Classification App") |
|
|
st.sidebar.title("Choose Task") |
|
|
task = st.sidebar.selectbox("Select a task", ["Image Classification", "Text Classification"]) |
|
|
|
|
|
if task == "Image Classification": |
|
|
st.header("Image Classification") |
|
|
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"]) |
|
|
if uploaded_file is not None: |
|
|
|
|
|
image = Image.open(uploaded_file) |
|
|
st.image(image, caption="Uploaded Image", use_column_width=True) |
|
|
|
|
|
|
|
|
if st.button("Classify Image"): |
|
|
with st.spinner("Classifying..."): |
|
|
results = image_classifier(image) |
|
|
st.subheader("Classification Results") |
|
|
for result in results: |
|
|
st.write(f"**{result['label']}**: {result['score']:.2f}") |
|
|
|
|
|
elif task == "Text Classification": |
|
|
st.header("Text Classification") |
|
|
text_input = st.text_area("Enter text for classification", "Streamlit is an amazing tool!") |
|
|
|
|
|
|
|
|
if st.button("Classify Text"): |
|
|
with st.spinner("Classifying..."): |
|
|
results = text_classifier(text_input) |
|
|
st.subheader("Classification Results") |
|
|
for result in results: |
|
|
st.write(f"**{result['label']}**: {result['score']:.2f}") |
|
|
|
|
|
st.write("Powered by Streamlit and Hugging Face 🤗") |
|
|
|