app.py

import streamlit as st
from transformers import (
    AutoModelForImageClassification, 
    AutoImageProcessor,
    ViTForImageClassification,
    ResNetForImageClassification
)
import torch
import numpy as np
from PIL import Image, ImageDraw
import cv2
from langchain import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain.llms import HuggingFacePipeline
import json
import os
from concurrent.futures import ThreadPoolExecutor
import pandas as pd

class DefectMeasurement:
    """Handle defect measurements and severity estimation"""
    
    @staticmethod
    def measure_defect(image, defect_type):
        """Measure defect dimensions using computer vision"""
        img_array = np.array(image)
        gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
        
        if defect_type == "Crack":
            # Crack width measurement
            blur = cv2.GaussianBlur(gray, (3,3), 0)
            edges = cv2.Canny(blur, 100, 200)
            lines = cv2.HoughLinesP(edges, 1, np.pi/180, 50, minLineLength=100, maxLineGap=10)
            
            if lines is not None:
                max_length = 0
                for line in lines:
                    x1, y1, x2, y2 = line[0]
                    length = np.sqrt((x2-x1)**2 + (y2-y1)**2)
                    max_length = max(max_length, length)
                return {"length": max_length, "unit": "pixels"}
            
        elif defect_type in ["Spalling", "Exposed_Bars"]:
            # Area measurement
            thresh = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)[1]
            contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            
            if contours:
                max_area = max(cv2.contourArea(cnt) for cnt in contours)
                return {"area": max_area, "unit": "square pixels"}
        
        return None

class MultiModelAnalyzer:
    """Handle multiple pre-trained models for defect detection"""
    
    def __init__(self):
        self.models = {
            "CODEBRIM-ViT": "chanwooong/codebrim-vit-base",
            "Concrete-Defect-ResNet": "nlp-waseda/concrete-defect-resnet",
            "Bridge-Damage-ViT": "microsoft/bridge-damage-vit-base"
        }
        self.loaded_models = {}
        self.loaded_processors = {}
        
    @st.cache_resource
    def load_model(self, model_name):
        """Load specific model and processor"""
        try:
            if "vit" in model_name.lower():
                model = ViTForImageClassification.from_pretrained(self.models[model_name])
            else:
                model = ResNetForImageClassification.from_pretrained(self.models[model_name])
            processor = AutoImageProcessor.from_pretrained(self.models[model_name])
            return model, processor
        except Exception as e:
            st.error(f"Error loading {model_name}: {str(e)}")
            return None, None
            
    def analyze_with_all_models(self, image):
        """Run analysis with all available models"""
        results = {}
        for model_name in self.models.keys():
            if model_name not in self.loaded_models:
                self.loaded_models[model_name], self.loaded_processors[model_name] = self.load_model(model_name)
            
            if self.loaded_models[model_name] is not None:
                try:
                    inputs = self.loaded_processors[model_name](images=image, return_tensors="pt")
                    outputs = self.loaded_models[model_name](**inputs)
                    probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
                    results[model_name] = probs
                except Exception as e:
                    st.error(f"Error analyzing with {model_name}: {str(e)}")
        
        return results

class EnhancedRAGSystem:
    """Enhanced RAG system with comprehensive construction knowledge"""
    
    def __init__(self):
        self.knowledge_sources = {
            "ACI_318": "concrete_design_requirements.json",
            "ASTM": "testing_standards.json",
            "repair_guidelines": "repair_methods.json",
            "case_studies": "defect_cases.json"
        }
        self.embeddings = None
        self.vectorstore = None
        self.qa_chain = None
        
    def load_knowledge_base(self):
        """Load and combine multiple knowledge sources"""
        combined_knowledge = []
        for source, filename in self.knowledge_sources.items():
            try:
                with open(f"knowledge_base/{filename}", 'r') as f:
                    knowledge = json.load(f)
                    for item in knowledge:
                        item['source'] = source
                    combined_knowledge.extend(knowledge)
            except Exception as e:
                st.warning(f"Could not load {source}: {str(e)}")
        
        return combined_knowledge
        
    def init_rag(self):
        """Initialize enhanced RAG system"""
        try:
            self.embeddings = HuggingFaceEmbeddings(
                model_name="sentence-transformers/all-mpnet-base-v2"
            )
            
            knowledge_base = self.load_knowledge_base()
            texts = [
                f"{item['defect_type']} ({item['source']})\n" +
                f"Description: {item['description']}\n" +
                f"Repair: {item['repair_methods']}\n" +
                f"Standards: {item['applicable_standards']}\n" +
                f"Cases: {item['related_cases']}"
                for item in knowledge_base
            ]
            
            self.vectorstore = FAISS.from_texts(texts, self.embeddings)
            
            self.qa_chain = RetrievalQA.from_chain_type(
                llm=HuggingFacePipeline.from_model_id(
                    model_id="google/flan-t5-large",
                    task="text2text-generation",
                    model_kwargs={"temperature": 0.7}
                ),
                chain_type="stuff",
                retriever=self.vectorstore.as_retriever(
                    search_kwargs={"k": 5}
                )
            )
            
            return True
        except Exception as e:
            st.error(f"Error initializing RAG system: {str(e)}")
            return False

class ConstructionDefectAnalyzer:
    """Main application class"""
    
    def __init__(self):
        self.multi_model = MultiModelAnalyzer()
        self.rag_system = EnhancedRAGSystem()
        self.defect_measurement = DefectMeasurement()
        
    def analyze_multiple_images(self, images):
        """Analyze multiple images in parallel"""
        results = []
        with ThreadPoolExecutor() as executor:
            futures = []
            for img in images:
                future = executor.submit(self.analyze_single_image, img)
                futures.append(future)
            
            for future in futures:
                result = future.result()
                results.append(result)
        
        return results
    
    def analyze_single_image(self, image):
        """Analyze a single image with all features"""
        model_results = self.multi_model.analyze_with_all_models(image)
        measurements = {}
        recommendations = {}
        
        # Get measurements for detected defects
        for model_name, predictions in model_results.items():
            for idx, prob in enumerate(predictions):
                if prob > 0.15:  # Confidence threshold
                    defect_type = self.get_defect_type(model_name, idx)
                    measurements[defect_type] = self.defect_measurement.measure_defect(image, defect_type)
                    
                    # Get RAG recommendations
                    if self.rag_system.qa_chain:
                        query = self.generate_rag_query(defect_type, measurements.get(defect_type))
                        recommendations[defect_type] = self.rag_system.qa_chain.run(query)
        
        return {
            "model_results": model_results,
            "measurements": measurements,
            "recommendations": recommendations
        }
    
    @staticmethod
    def generate_rag_query(defect_type, measurement):
        """Generate detailed query for RAG system"""
        query = f"What are the recommended repairs, safety measures, and applicable standards for {defect_type}"
        if measurement:
            if "length" in measurement:
                query += f" with length {measurement['length']} {measurement['unit']}"
            elif "area" in measurement:
                query += f" with affected area {measurement['area']} {measurement['unit']}"
        return query + "?"

def main():
    st.set_page_config(page_title="Advanced Construction Defect Analyzer", layout="wide")
    
    analyzer = ConstructionDefectAnalyzer()
    
    st.title("🏗️ Advanced Construction Defect Analyzer")
    
    # Multiple image upload
    uploaded_files = st.file_uploader(
        "Upload construction images for analysis",
        type=['jpg', 'jpeg', 'png'],
        accept_multiple_files=True
    )
    
    if uploaded_files:
        images = [Image.open(file).convert('RGB') for file in uploaded_files]
        
        with st.spinner("Analyzing images..."):
            results = analyzer.analyze_multiple_images(images)
            
            for idx, (image, result) in enumerate(zip(images, results)):
                st.markdown(f"### Analysis Results - Image {idx + 1}")
                
                col1, col2 = st.columns([1, 2])
                
                with col1:
                    st.image(image, caption=f"Image {idx + 1}", use_column_width=True)
                
                with col2:
                    # Display model comparison
                    st.markdown("#### Model Predictions")
                    for model_name, predictions in result['model_results'].items():
                        st.markdown(f"**{model_name}:**")
                        for i, prob in enumerate(predictions):
                            if prob > 0.15:
                                defect_type = analyzer.get_defect_type(model_name, i)
                                st.progress(float(prob))
                                st.markdown(f"{defect_type}: {float(prob)*100:.1f}%")
                                
                                # Display measurements
                                if defect_type in result['measurements']:
                                    st.markdown("**Measurements:**")
                                    st.json(result['measurements'][defect_type])
                                
                                # Display recommendations
                                if defect_type in result['recommendations']:
                                    with st.expander("📋 Detailed Analysis"):
                                        st.markdown(result['recommendations'][defect_type])

if __name__ == "__main__":
    main()