Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import streamlit as st | |
| from PIL import Image | |
| import torch | |
| import torch.nn as nn | |
| import torchvision.transforms as transforms | |
| from utils import preprocess_image | |
| import torch | |
| from transformers import ViTForImageClassification | |
| labels = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash'] | |
| # Initialize a ViT model | |
| model = ViTForImageClassification.from_pretrained("facebook/deit-base-distilled-patch16-224") | |
| # Load your weights | |
| model.load_state_dict(torch.load("best.pt", map_location="cpu")) | |
| model.eval() # Set to evaluation mode | |
| # Customized Streamlit layout | |
| st.set_page_config( | |
| page_title="EcoIdentify by EcoClim Solutions", | |
| page_icon="https://ecoclimsolutions.files.wordpress.com/2024/01/rmcai-removebg.png?resize=48%2C48", | |
| layout="wide", | |
| initial_sidebar_state="expanded", | |
| ) | |
| # Customized Streamlit styles | |
| st.markdown( | |
| """ | |
| <style> | |
| body { | |
| color: #333333; | |
| background-color: #f9f9f9; | |
| font-family: 'Helvetica', sans-serif; | |
| } | |
| .st-bb { | |
| padding: 0rem; | |
| } | |
| .st-ec { | |
| color: #666666; | |
| } | |
| .st-ef { | |
| color: #666666; | |
| } | |
| .st-ei { | |
| color: #333333; | |
| } | |
| .st-dh { | |
| font-size: 36px; | |
| font-weight: bold; | |
| color: #4CAF50; | |
| text-align: center; | |
| margin-bottom: 20px; | |
| } | |
| .st-gf { | |
| background-color: #4CAF50; | |
| color: white; | |
| padding: 15px 30px; | |
| font-size: 18px; | |
| border: none; | |
| border-radius: 8px; | |
| cursor: pointer; | |
| transition: background-color 0.3s; | |
| } | |
| .st-gf:hover { | |
| background-color: #45a049; | |
| } | |
| .st-gh { | |
| text-align: center; | |
| font-size: 24px; | |
| font-weight: bold; | |
| margin-bottom: 20px; | |
| } | |
| .st-logo { | |
| max-width: 100%; | |
| height: auto; | |
| margin: 20px auto; | |
| display: block; | |
| } | |
| </style> | |
| """, | |
| unsafe_allow_html=True, | |
| ) | |
| # Logo | |
| st.image("https://ecoclimsolutions.files.wordpress.com/2024/01/rmcai-removebg.png?resize=48%2C48") | |
| # Page title | |
| st.title("EcoIdentify by EcoClim Solutions") | |
| # Subheader | |
| st.header("Upload a waste image to find its category") | |
| # Note | |
| st.markdown("* Please note that our dataset is trained primarily with images that contain a white background. Therefore, images with white background would produce maximum accuracy *") | |
| # Image upload section | |
| opt = st.selectbox("How do you want to upload the image for classification?", ("Please Select", "Upload image from device")) | |
| image = None | |
| if opt == 'Upload image from device': | |
| file = st.file_uploader('Select', type=['jpg', 'png', 'jpeg']) | |
| if file: | |
| image = preprocess_image(file) | |
| try: | |
| if image is not None: | |
| st.image(image, width=256, caption='Uploaded Image') | |
| if st.button('Predict'): | |
| transform = transforms.Compose([ | |
| transforms.Resize((256, 256)), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), | |
| ]) | |
| image = transform(image).unsqueeze(0) | |
| with torch.no_grad(): | |
| prediction = model(image) | |
| st.success(f'Prediction: {labels[torch.argmax(prediction, dim=1).item()]}') | |
| except Exception as e: | |
| st.error(f"An error occurred: {e}. Please contact us EcoClim Solutions at EcoClimSolutions.wordpress.com.") | |