Update app.py

app.py

@@ -1,87 +1,38 @@
 import streamlit as st
 from PIL import Image
 import torch
-from transformers import pipeline
-import time
+from transformers import pipeline, AutoFeatureExtractor
 
-# Custom CSS for styling
-st.markdown("""
-    <style>
-    .title {
-        font-size: 40px;
-        color: #4CAF50;
-        font-weight: bold;
-        text-align: center;
-    }
-    .sub-title {
-        font-size: 20px;
-        color: #333;
-        text-align: center;
-    }
-    .upload-button {
-        background-color: #4CAF50;
-        color: white;
-        border-radius: 5px;
-        padding: 10px 20px;
-        font-size: 18px;
-    }
-    .upload-button:hover {
-        background-color: #45a049;
-    }
-    .prediction {
-        font-size: 22px;
-        font-weight: bold;
-        color: #388E3C;
-        text-align: center;
-    }
-    .confidence {
-        font-size: 18px;
-        color: #555;
-        text-align: center;
-    }
-    .image-container {
-        display: flex;
-        justify-content: center;
-        margin-top: 30px;
-    }
-    .footer {
-        text-align: center;
-        margin-top: 50px;
-        font-size: 14px;
-        color: #777;
-    }
-    </style>
-    """, unsafe_allow_html=True)
-
-# Load model from Hugging Face
+# Load the model and feature extractor (image processor)
 @st.cache_resource
 def load_model():
     model = pipeline('image-classification', model='google/vit-base-patch16-224-in21k')
     return model
 
+# Initialize the feature extractor
+extractor = AutoFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
+
 model = load_model()
 
 # Streamlit app UI
-st.markdown('<p class="title">🌱 Plant Identification App 🌱</p>', unsafe_allow_html=True)
-st.markdown('<p class="sub-title">Upload a plant image and let the app identify its species!</p>', unsafe_allow_html=True)
+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"], label_visibility="collapsed")
+uploaded_file = st.file_uploader("Choose a plant image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
-    # Show the uploaded image with a loading effect
+    # Open and display the uploaded image
     image = Image.open(uploaded_file)
     st.image(image, caption="Uploaded Plant Image.", use_column_width=True)
 
-    # Display loading indicator and classify
-    with st.spinner('Classifying plant species...'):
-        time.sleep(2)  # Simulate loading time
-        predictions = model(image)
+    # Preprocess the image with the feature extractor
+    inputs = extractor(images=image, return_tensors="pt", padding=True)
 
-    # Display prediction result
-    st.markdown(f'<p class="prediction">Predicted Species: {predictions[0]["label"]}</p>', unsafe_allow_html=True)
-    st.markdown(f'<p class="confidence">Confidence: {predictions[0]["score"]*100:.2f}%</p>', unsafe_allow_html=True)
-
-# Footer for extra info or tips
-st.markdown('<div class="footer">Powered by Streamlit 🌸</div>', unsafe_allow_html=True)
+    # Run the classification
+    with st.spinner('Classifying plant species...'):
+        predictions = model(inputs['pixel_values'])
 
+    # Display prediction results
+    st.write(f"Predicted Species: {predictions[0]['label']}")
+    st.write(f"Confidence: {predictions[0]['score']*100:.2f}%")