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

# Load the model
learn = load_learner('model.pkl')

st.set_page_config(page_title="Breast Cancer Detector", layout="centered")
st.title("🩺 Breast Cancer Detection (MIAS ROI Dataset)")

st.write("Upload a mammogram ROI image to see if it's **cancerous** or **normal**.")

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

# Select input source
img = None
if uploaded_file is not None:
    img = Image.open(uploaded_file).convert('RGB')

if img is not None:
    st.image(img, caption='Input Image', use_column_width=True)

    # Convert to fastai's PILImage
    img_fastai = PILImage.create(img)

    # Prediction
    pred_class, pred_idx, probs = learn.predict(img_fastai)

    # Show results
    st.subheader(f"Prediction: {pred_class.upper()}")
    
    # Show both class probabilities
    for cls, prob in zip(learn.dls.vocab, probs):
        st.write(f"**{cls.capitalize()}**: {prob*100:.2f}% confidence")

st.markdown("---")

# Sample Images
st.subheader("Try with Sample Images:")

col1, col2 = st.columns(2)

with col1:
    if st.button('Sample Normal'):
        sample_path = 'normal_sample.png'  # You must upload this in Hugging Face
        img = Image.open(sample_path).convert('RGB')
        st.image(img, caption="Normal Sample", use_column_width=True)
        img_fastai = PILImage.create(img)
        pred_class, pred_idx, probs = learn.predict(img_fastai)
        st.subheader(f"Prediction: {pred_class.upper()}")
        for cls, prob in zip(learn.dls.vocab, probs):
            st.write(f"**{cls.capitalize()}**: {prob*100:.2f}% confidence")

with col2:
    if st.button('Sample Cancer'):
        sample_path = 'cancer_sample.png'  # You must upload this too
        img = Image.open(sample_path).convert('RGB')
        st.image(img, caption="Cancer Sample", use_column_width=True)
        img_fastai = PILImage.create(img)
        pred_class, pred_idx, probs = learn.predict(img_fastai)
        st.subheader(f"Prediction: {pred_class.upper()}")
        for cls, prob in zip(learn.dls.vocab, probs):
            st.write(f"**{cls.capitalize()}**: {prob*100:.2f}% confidence")

st.markdown("---")
st.write("Made by Novoo.")