import streamlit as st
from transformers import pipeline, AutoModelForImageClassification, AutoImageProcessor
from PIL import Image
import torch

# Initialize models
model_name = "microsoft/resnet-50"
image_processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)
classifier = pipeline("image-classification", model=model, image_processor=image_processor)

damage_types = {
    'spalling': ['Remove loose material', 'Apply repair mortar', 'Seal surface'],
    'cracking': ['Measure crack width', 'Install crack monitors', 'Apply sealant'],
    'corrosion': ['Clean exposed steel', 'Apply rust converter', 'Paint protective coating'],
    'dampness': ['Locate water source', 'Fix drainage issues', 'Apply waterproofing']
}

def analyze_image(image):
    return classifier(image)

def main():
    st.title("Structural Damage Assessment")
    
    uploaded_file = st.file_uploader("Upload structural image", type=['jpg', 'jpeg', 'png'])
    
    if uploaded_file:
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image")
        
        with st.spinner("Analyzing..."):
            results = analyze_image(image)
            
            for result in results:
                confidence = result['score'] * 100
                st.progress(confidence / 100)
                st.write(f"{result['label']}: {confidence:.1f}%")
                
                for damage_type, remedies in damage_types.items():
                    if damage_type in result['label'].lower():
                        st.subheader("Recommended Actions:")
                        for remedy in remedies:
                            st.write(f"• {remedy}")

if __name__ == "__main__":
    main()