Update app.py

app.py

@@ -5,12 +5,16 @@ import torch
 import torch.nn as nn
 import torchvision.transforms as transforms
 from utils import preprocess_image
+import torch
+from transformers import ViTForImageClassification
 
-# Initialize labels and model
 labels = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
-model = torch.load('classifier.h5')
-model.eval()
+# 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",