app.py

import streamlit as st
import torch
from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
import numpy as np
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
import io
import json

# Load pre-trained model and processor
model_name = "google/vit-base-patch16-224-in21k"
processor = ViTImageProcessor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)

# Custom class labels
damage_types = [
    "spalling",
    "reinforcement_corrosion",
    "flexural_crack",
    "structural_crack",
    "dampness",
    "impact_failure"
]

# Initialize FAISS vector store
embeddings = HuggingFaceEmbeddings()
knowledge_base = FAISS.load_local("knowledge_base", embeddings)

def process_image(image):
    # Preprocess image
    inputs = processor(images=image, return_tensors="pt")
    
    # Get model predictions
    outputs = model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    
    # Get top predictions
    top_probs, top_indices = torch.topk(probs, len(damage_types))
    
    return {
        damage_types[idx]: float(prob)
        for idx, prob in zip(top_indices[0], top_probs[0])
    }

def get_recommendations(damage_type):
    # Query vector store for recommendations
    docs = knowledge_base.similarity_search(
        f"Remedial measures for {damage_type} in building structures",
        k=3
    )
    return [doc.page_content for doc in docs]

# Streamlit UI
st.title("Structural Damage Assessment Tool")

# File upload
uploaded_file = st.file_uploader("Upload structural image", type=["jpg", "jpeg", "png"])

if uploaded_file:
    # Display image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)
    
    # Process image
    with st.spinner("Analyzing image..."):
        predictions = process_image(image)
    
    # Display results
    st.subheader("Damage Assessment")
    for damage_type, probability in predictions.items():
        st.progress(probability)
        st.write(f"{damage_type.replace('_', ' ').title()}: {probability:.2%}")
        
        # Show recommendations
        if probability > 0.5:
            st.subheader(f"Recommendations for {damage_type.replace('_', ' ').title()}")
            recommendations = get_recommendations(damage_type)
            for i, rec in enumerate(recommendations, 1):
                st.write(f"{i}. {rec}")
    
    # Generate report
    st.download_button(
        "Download Report",
        json.dumps(predictions, indent=2),
        "assessment_report.json",
        "application/json"
    )