MahdiHasan/Image_Classifier

1import streamlit as st
2import torch
3import numpy as np
4from torchvision import transforms
5from PIL import Image
6from models import ResNet
7
8# Class labels
9class_names = ['berry', 'bird', 'dog', 'flower']
10
11# Load the saved model
12def load_model(model, filename):
13    model.load_state_dict(torch.load(filename, map_location=torch.device('cpu')))
14    model.eval()  # Set the model to evaluation mode
15    return model
16
17preprocess = transforms.Compose([
18    transforms.Resize((224, 224)), 
19    transforms.ToTensor()
20])
21
22# Inference function
23def predict(image, model, class_names):
24    image = preprocess(image)
25    image = image.unsqueeze(0) 
26    
27    with torch.no_grad():
28        output = model(image)
29    
30    _, predicted_class = torch.max(output, 1)
31    predicted_label = class_names[predicted_class.item()]
32    return predicted_label
33
34# Load the model
35model = ResNet(num_classes=4)
36model = load_model(model, 'resnet_model.pth')
37
38# Streamlit app
39st.title("Image Classification")
40st.write("Upload an image to predict the label.")
41
42uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
43
44if uploaded_file is not None:
45    # Display uploaded image
46    img = Image.open(uploaded_file)
47    st.image(img, caption="Uploaded Image", use_container_width=True)
48    
49    # Predict font
50    predicted_label = predict(img, model, class_names)
51    st.write(f"Predicted Label: **{predicted_label}**")
52