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

# 1. Setup the Page
st.title("Concrete Crack Detection")
st.write("Upload an image of a concrete surface, and the AI will predict if there is a crack.")

# 2. Load the Model 
# Use a public model for now since you don't have a trained one yet
# This model is a common one for crack detection
MODEL_ID = "Guna-M/concrete-crack-ai" 

try:
    classifier = pipeline("image-classification", model=MODEL_ID)
except Exception as e:
    st.error(f"Error loading model: {e}")
    st.stop()

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

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image', use_container_width=True)
    st.write("Analyzing...")
    
    # 4. Make Prediction
    results = classifier(image)
    
    st.subheader("Prediction Results:")
    
    # Show the top result
    top_result = results[0]
    label = top_result['label']
    score = top_result['score'] * 100
    
    if "crack" in label.lower() and "no" not in label.lower():
        st.error(f"Prediction: CRACK DETECTED (Confidence: {score:.2f}%)")
    else:
        st.success(f"Prediction: NO CRACK (Confidence: {score:.2f}%)")

    # Show all probabilities
    with st.expander("See all probabilities"):
        for res in results:
            st.write(f"- {res['label']}: {res['score']*100:.2f}%")