import subprocess
import sys

def install(package):
    subprocess.check_call([sys.executable, "-m", "pip", "install", package])

install('fastai')
install('torch')

import streamlit as st
from fastai.vision.all import *
from PIL import Image

# Load the model (ensure model file is in the same directory or provide correct path)
@st.cache(allow_output_mutation=True)
def load_model():
    # Replace 'model.pkl' with your model file
    model = load_learner('modelmarvel.pkl')
    return model

# Load the model
model = load_model()

# Streamlit application
st.title("Image Classification with fastai")

# Image upload
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])

if uploaded_file:
    # Open and display the image
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded Image", use_column_width=True)
    
    # Prepare image for the model
    img = PILImage.create(uploaded_file)
    
    # Predict
    pred, pred_idx, probs = model.predict(img)
    
    # Display results
    st.write(f"Prediction: {pred}")
    st.write(f"Probability: {probs[pred_idx]:.4f}")

    # Optionally, display confidence as a bar chart
    st.bar_chart(probs.numpy())