talha2001/Efficientnet

1import streamlit as st2import torch3import torchvision.transforms as T4from PIL import Image5from torchvision import models6import io7 8# Load the model9@st.cache_resource10def load_model():11    model = models.efficientnet_b0(pretrained=False)12    num_classes = 3  # Replace with the number of your classes13    model.classifier = torch.nn.Sequential(14        torch.nn.Dropout(p=0.2, inplace=True),15        torch.nn.Linear(in_features=1280, out_features=num_classes, bias=True)16    )17    model.load_state_dict(torch.load("efficientnet_b0_best_2nd.pth", map_location=torch.device('cpu')))18    model.eval()19    return model20 21model = load_model()22 23# Define the image transformation24def transform_image(image):25    transform = T.Compose([26        T.Resize((512, 512)),27        T.ToTensor(),28        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # ImageNet normalization29    ])30    return transform(image).unsqueeze(0)31 32# Function to predict the class33def predict_image(image):34    tensor = transform_image(image)35    with torch.no_grad():36        output = model(tensor)37    _, predicted_class = torch.max(output, 1)38    return predicted_class.item()39 40# Define class labels (modify based on your classes)41class_names = ['ACNE', 'Eczema', 'Psoriasis']42 43# Streamlit interface44st.title("Image Classification with EfficientNet")45st.write("Upload an image to classify it.")46 47# File uploader48uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])49 50if uploaded_file is not None:51    # Open and display the uploaded image52    image = Image.open(uploaded_file)53    st.image(image, caption="Uploaded Image.", use_column_width=True)54 55    # Predict and display the result56    predicted_class = predict_image(image)57    st.write(f"Prediction: {class_names[predicted_class]}")58 59