import streamlit as st
from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
import torch

# Custom labels for structural damage
DAMAGE_TYPES = {
    'spalling': {
        'remedies': ['Remove loose concrete', 'Clean exposed reinforcement', 'Apply repair mortar'],
        'risk_level': 'High'
    },
    'reinforcement_corrosion': {
        'remedies': ['Remove rust', 'Apply corrosion inhibitor', 'Repair concrete cover'],
        'risk_level': 'Critical'
    },
    'structural_crack': {
        'remedies': ['Monitor crack width', 'Epoxy injection', 'Structural reinforcement'],
        'risk_level': 'Critical'
    },
    'dampness': {
        'remedies': ['Identify water source', 'Fix waterproofing', 'Improve drainage'],
        'risk_level': 'Medium'
    }
}

@st.cache_resource
def load_model():
    model = ViTForImageClassification.from_pretrained(
        "nateraw/vit-base-patch16-224-finetuned-structural-damage",
        num_labels=len(DAMAGE_TYPES),
        id2label={i: label for i, label in enumerate(DAMAGE_TYPES.keys())},
        label2id={label: i for i, label in enumerate(DAMAGE_TYPES.keys())}
    )
    processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
    return model, processor

def analyze_damage(image, model, processor):
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    return probs[0]

def main():
    st.title("Structural Damage Assessment")
    
    model, processor = load_model()
    
    uploaded_file = st.file_uploader("Upload structural image", type=['jpg', 'jpeg', 'png'])
    
    if uploaded_file:
        image = Image.open(uploaded_file)
        st.image(image, caption="Analyzed Structure")
        
        with st.spinner("Analyzing structural damage..."):
            predictions = analyze_damage(image, model, processor)
            
            for idx, (damage_type, conf) in enumerate(zip(DAMAGE_TYPES.keys(), predictions)):
                confidence = float(conf) * 100
                if confidence > 10:  # Show only significant predictions
                    st.subheader(f"{damage_type.replace('_', ' ').title()}")
                    st.progress(confidence / 100)
                    st.write(f"Confidence: {confidence:.1f}%")
                    st.write(f"Risk Level: {DAMAGE_TYPES[damage_type]['risk_level']}")
                    
                    st.write("Recommended Actions:")
                    for remedy in DAMAGE_TYPES[damage_type]['remedies']:
                        st.write(f"• {remedy}")
                    st.divider()

if __name__ == "__main__":
    main()