File size: 1,490 Bytes
de1be37 7d493ea de1be37 7d493ea de1be37 7d493ea de1be37 7d493ea de1be37 bce5713 de1be37 bce5713 de1be37 bce5713 de1be37 7d493ea de1be37 bce5713 7d493ea de1be37 bce5713 7d493ea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import streamlit as st
from PIL import Image
from transformers import AutoModelForImageClassification, ViTImageProcessor
import torch
# Load the model and processor
@st.cache_resource
def load_model():
model = AutoModelForImageClassification.from_pretrained("google/vit-base-patch16-224-in21k")
processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
return model, processor
model, processor = load_model()
# Streamlit app UI
st.title("🌱 Plant Identification App 🌱")
st.write("Upload a plant image and let the app identify its species!")
# File uploader for plant image
uploaded_file = st.file_uploader("Choose a plant image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
# Open and display the uploaded image
image = Image.open(uploaded_file)
st.image(image, caption="Uploaded Plant Image.", use_column_width=True)
# Preprocess the image using the processor
inputs = processor(images=image, return_tensors="pt", padding=True)
# Run the classification
with st.spinner('Classifying plant species...'):
outputs = model(**inputs)
logits = outputs.logits
predicted_class_idx = logits.argmax(-1).item()
# Get the label of the predicted class
label = model.config.id2label[predicted_class_idx]
# Display prediction results
st.write(f"Predicted Species: {label}")
st.write(f"Confidence: {torch.softmax(logits, dim=-1)[0][predicted_class_idx]*100:.2f}%")
|