Update app.py

app.py

@@ -1,79 +1,273 @@
 import streamlit as st
-from transformers import ViTForImageClassification, ViTImageProcessor
-from PIL import Image
+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
 
-# Define damage types and remedies
-DAMAGE_TYPES = {
-    0: {'name': 'spalling', 'risk': 'High'},
-    1: {'name': 'reinforcement_corrosion', 'risk': 'Critical'},
-    2: {'name': 'structural_crack', 'risk': 'High'},
-    3: {'name': 'dampness', 'risk': 'Medium'},
-    4: {'name': 'no_damage', 'risk': 'Low'}
-}
+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
 
-REMEDIES = {
-    'spalling': ['Remove loose concrete', 'Clean exposed area', 'Apply repair mortar'],
-    'reinforcement_corrosion': ['Remove rust', 'Apply corrosion inhibitor', 'Repair concrete cover'],
-    'structural_crack': ['Measure crack width', 'Epoxy injection', 'Monitor progression'],
-    'dampness': ['Identify water source', 'Improve drainage', 'Apply waterproofing'],
-    'no_damage': ['Regular maintenance', 'Periodic inspection']
-}
+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
 
-@st.cache_resource
-def load_model():
-    model = ViTForImageClassification.from_pretrained(
-        "google/vit-base-patch16-224",
-        num_labels=len(DAMAGE_TYPES),
-        ignore_mismatched_sizes=True
-    )
-    processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
-    return model, processor
+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
 
-def analyze_damage(image, model, processor):
-    image = image.convert('RGB')
-    inputs = processor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
-    probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
-    return probs
+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.title("Structural Damage Assessment Tool")
-
-    model, processor = load_model()
+    st.set_page_config(page_title="Advanced Construction Defect Analyzer", layout="wide")
+    
+    analyzer = ConstructionDefectAnalyzer()
+    
+    st.title("🏗️ Advanced Construction Defect Analyzer")
     
-    uploaded_file = st.file_uploader("Upload structural image", type=['jpg', 'jpeg', 'png'])
+    # Multiple image upload
+    uploaded_files = st.file_uploader(
+        "Upload construction images for analysis",
+        type=['jpg', 'jpeg', 'png'],
+        accept_multiple_files=True
+    )
     
-    if uploaded_file:
-        image = Image.open(uploaded_file)
-        st.image(image, caption="Uploaded Structure", use_column_width=True)
+    if uploaded_files:
+        images = [Image.open(file).convert('RGB') for file in uploaded_files]
         
-        with st.spinner("Analyzing..."):
-            predictions = analyze_damage(image, model, processor)
-            
-            col1, col2 = st.columns(2)
+        with st.spinner("Analyzing images..."):
+            results = analyzer.analyze_multiple_images(images)
             
-            with col1:
-                st.subheader("Damage Assessment")
-                for idx, prob in enumerate(predictions):
-                    damage_type = DAMAGE_TYPES[idx]['name']
-                    confidence = float(prob) * 100
-                    if confidence > 15:
-                        st.write(f"**{damage_type.replace('_', ' ').title()}**")
-                        st.progress(confidence / 100)
-                        st.write(f"Confidence: {confidence:.1f}%")
-                        st.write(f"Risk Level: {DAMAGE_TYPES[idx]['risk']}")
-            
-            with col2:
-                st.subheader("Recommended Actions")
-                for idx, prob in enumerate(predictions):
-                    damage_type = DAMAGE_TYPES[idx]['name']
-                    confidence = float(prob) * 100
-                    if confidence > 15:
-                        st.write(f"**For {damage_type.replace('_', ' ').title()}:**")
-                        for remedy in REMEDIES[damage_type]:
-                            st.write(f"• {remedy}")
-                        st.write("---")
+            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()
+    main()
+