import streamlit as st
import torch
import torchvision.transforms as transforms
from PIL import Image
import numpy as np

# Load your pre-trained model
model = torch.load('model/your_model_file.pt')
model.eval()

# Define image transformations
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

st.title("VIEP: Utility Pole Fault Detection")

uploaded_file = st.file_uploader("Upload an image of a utility pole", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    image = Image.open(uploaded_file).convert('RGB')
    st.image(image, caption='Uploaded Image', use_column_width=True)

    # Preprocess the image
    input_tensor = transform(image).unsqueeze(0)

    # Perform inference
    with torch.no_grad():
        output = model(input_tensor)
        _, predicted = torch.max(output, 1)

    # Map the prediction to class names
    classes = ['No Fault', 'Fault Detected']
    st.write(f"Prediction: {classes[predicted.item()]}")