import streamlit as st
from transformers import pipeline
from PIL import Image

# Load the pre-trained image classification pipeline
pipe = pipeline("image-classification", model="ALM-AHME/convnextv2-large-1k-224-finetuned-BreastCancer-Classification-BreakHis-AH-60-20-20")

# Define the decision logic based on confidence scores
def classify_mammogram(img):
    results = pipe(img)
    predicted_label = results[0]['label']
    confidence = results[0]['score']
    
    if predicted_label == "malignant":
        if confidence >= 0.8:
            classification = "Non-suspicious"
        elif 0.6 <= confidence < 0.8:
            classification = "High risk"
        else:
            classification = "Indeterminate"
    else:  # Benign prediction
        if confidence < 0.4:
            classification = "Suspicious"
        elif 0.4 <= confidence < 0.8:
            classification = "Non-suspicious"
        else:
            classification = "No risk"

    return f"Predicted Class: {predicted_label}\nClassification: {classification}\nConfidence: {confidence:.2f}"

# Streamlit app
st.title("Breast Cancer Detection App")
st.write("Upload an image of a mammogram(an X-ray image of the breast), and the model will predict whether it is benign or malignant.")


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

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image.', use_column_width=True)
    st.write("Classifying...")
    
    classification_result = classify_mammogram(image)
    st.write(classification_result)

st.write(
    """
    <div style='display: flex; justify-content: space-between;'>
        <div style='width: 48%;'>
            <p style='font-size: 18px; font-weight: bold;'>Malignant</p>
            <p style='font-size: 16px;'>
                Definition: very virulent or infectious.
                Non-suspicious: The AI is confident that no suspicious signs are present.<br>
                High risk: The AI is confident that the results are highly suspicious.<br>
                Indeterminate: The AI is uncertain and not confident in making a definitive classification.
            </p>
        </div>
        <div style='width: 48%;'>
            <p style='font-size: 18px; font-weight: bold;'>Benign</p>
            <p style='font-size: 16px;'>
                Definition: Not harmful in effect
                Suspicious: AI is not that confident. Further Supervision is needed.<br>
                Non-Suspicious: AI is confident that there is nothing to worry about.<br>
                No Risk: AI is confident. No direct supervision is needed.
            </p>
        </div>
    </div>
    """, 
    unsafe_allow_html=True
)

# Educational content
st.markdown("### Learn More About Breast Cancer")
st.markdown("""
    - [Breast Cancer Overview](https://www.who.int/news-room/fact-sheets/detail/breast-cancer)
    - [Preventing Breast Cancer](https://www.cancer.gov/types/breast/patient/breast-prevention-pdq)
    """)
# Determine the current Streamlit theme (light or dark)
theme = st.get_option("theme.base")

    # Define button styling based on theme
if theme == "light":
    button_bg_color = "#2c2e35"
    button_border_color = "1px solid black"
    button_text_color = "black"
else:
    button_bg_color = "#2c2e35"
    button_border_color = "1px solid #fff"
    button_text_color = "#fff"

    # Rounded button-like element with dynamic styling
st.markdown(f"""
<style>
.rounded-button {{
    display: inline-block;
    padding: 7px 15px;
    font-size: 16px;
    color: {button_text_color};
    background-color: {button_bg_color};
    border: {button_border_color};
    border-radius: 7px;
    text-align: center;
    text-decoration: none;
    cursor: default;
}}
</style>
<div style="text-align: center;">
    <div class="rounded-button">
        Created by: Samuel Ameyaw
     </div>
    </div>
    """, unsafe_allow_html=True)