Update app.py

app.py

@@ -1,273 +1,325 @@
+```python
 import streamlit as st
-from transformers import (
-    AutoModelForImageClassification, 
-    AutoImageProcessor,
-    ViTForImageClassification,
-    ResNetForImageClassification
-)
+from transformers import ViTForImageClassification, ViTImageProcessor
+from PIL import Image
 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
+import time
+import gc
+from knowledge_base import KNOWLEDGE_BASE, DAMAGE_TYPES
+from rag_utils import RAGSystem
 
-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"}
-        
+# Constants
+MAX_FILE_SIZE = 5 * 1024 * 1024  # 5MB
+MAX_IMAGE_SIZE = 1024  # Maximum dimension for images
+
+# Cache the model and RAG system globally
+MODEL = None
+PROCESSOR = None
+RAG_SYSTEM = None
+
+def cleanup_memory():
+    """Clean up memory and GPU cache"""
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+def init_session_state():
+    """Initialize session state variables"""
+    if 'history' not in st.session_state:
+        st.session_state.history = []
+    if 'dark_mode' not in st.session_state:
+        st.session_state.dark_mode = False
+
+@st.cache_resource(show_spinner="Loading AI model...")
+def load_model():
+    """Load and cache the model and processor"""
+    try:
+        model_name = "google/vit-base-patch16-224"
+        model = ViTForImageClassification.from_pretrained(
+            model_name,
+            num_labels=len(DAMAGE_TYPES),
+            ignore_mismatched_sizes=True,
+            device_map="auto"
+        )
+        processor = ViTImageProcessor.from_pretrained(model_name)
+        return model, processor
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        return None, None
+
+def init_rag_system():
+    """Initialize the RAG system with knowledge base"""
+    global RAG_SYSTEM
+    if RAG_SYSTEM is None:
+        RAG_SYSTEM = RAGSystem()
+        RAG_SYSTEM.initialize_knowledge_base(KNOWLEDGE_BASE)
+
+def validate_image(image):
+    """Validate image size and format"""
+    if image.size[0] * image.size[1] > 1024 * 1024:
+        st.warning("Large image detected. The image will be resized for better performance.")
+    if image.format not in ['JPEG', 'PNG']:
+        st.warning("Image format not optimal. Consider using JPEG or PNG for better performance.")
+
+def preprocess_image(uploaded_file):
+    """Preprocess and validate uploaded image"""
+    try:
+        image = Image.open(uploaded_file)
+        # Resize if image is too large
+        if max(image.size) > MAX_IMAGE_SIZE:
+            ratio = MAX_IMAGE_SIZE / max(image.size)
+            new_size = tuple([int(dim * ratio) for dim in image.size])
+            image = image.resize(new_size, Image.Resampling.LANCZOS)
+        return image
+    except Exception as e:
+        st.error(f"Error processing image: {str(e)}")
         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
+def analyze_damage(image, model, processor):
+    """Analyze structural damage in the image"""
+    try:
+        with torch.no_grad():
+            image = image.convert('RGB')
+            inputs = processor(images=image, return_tensors="pt")
+            outputs = model(**inputs)
+            probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
+            cleanup_memory()
+            return probs
+    except RuntimeError as e:
+        if "out of memory" in str(e):
+            cleanup_memory()
+            st.error("Out of memory. Please try with a smaller image.")
+        else:
+            st.error(f"Error analyzing image: {str(e)}")
+        return None
 
-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 get_custom_css():
+    """Return custom CSS styles"""
+    return """
+    <style>
+    .main {
+        padding: 2rem;
+    }
+    .stProgress > div > div > div > div {
+        background-image: linear-gradient(to right, var(--progress-color, #ff6b6b), var(--progress-color-end, #f06595));
+    }
+    .damage-card {
+        padding: 1.5rem;
+        border-radius: 0.5rem;
+        background: var(--card-bg, #f8f9fa);
+        margin-bottom: 1rem;
+        border: 1px solid var(--border-color, #dee2e6);
+        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+    }
+    .damage-header {
+        font-size: 1.25rem;
+        font-weight: bold;
+        margin-bottom: 1rem;
+        color: var(--text-color, #212529);
+    }
+    .dark-mode {
+        background-color: #1a1a1a;
+        color: #ffffff;
+    }
+    .dark-mode .damage-card {
+        background: #2d2d2d;
+        border-color: #404040;
+    }
+    </style>
+    """
+
+def display_header():
+    """Display application header"""
+    st.markdown(
+        """
+        <div style='text-align: center; padding: 1rem;'>
+            <h1>🏗️ Structural Damage Analyzer Pro</h1>
+            <p style='font-size: 1.2rem;'>Advanced AI-powered structural damage assessment tool</p>
+        </div>
+        """,
+        unsafe_allow_html=True
+    )
+
+def display_enhanced_analysis(damage_type, confidence):
+    """Display enhanced analysis from RAG system"""
+    try:
+        enhanced_info = RAG_SYSTEM.get_enhanced_analysis(damage_type, confidence)
         
-    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)}")
+        st.markdown("### 🔍 Enhanced Analysis")
         
-        return combined_knowledge
+        with st.expander("📚 Technical Details", expanded=True):
+            for detail in enhanced_info["technical_details"]:
+                st.markdown(detail)
+                
+        with st.expander("⚠️ Safety Considerations"):
+            for safety in enhanced_info["safety_considerations"]:
+                st.warning(safety)
+                
+        with st.expander("👷 Expert Recommendations"):
+            for rec in enhanced_info["expert_recommendations"]:
+                st.info(rec)
+                
+        custom_query = st.text_input(
+            "Ask specific questions about this damage type:",
+            placeholder="E.g., What are the long-term implications of this damage?"
+        )
         
-    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}
-                )
+        if custom_query:
+            custom_results = RAG_SYSTEM.get_enhanced_analysis(
+                damage_type, 
+                confidence, 
+                custom_query=custom_query
             )
-            
-            return True
-        except Exception as e:
-            st.error(f"Error initializing RAG system: {str(e)}")
-            return False
+            st.markdown("### 💡 Custom Query Results")
+            for category, results in custom_results.items():
+                if results:
+                    st.markdown(f"**{category.replace('_', ' ').title()}:**")
+                    for result in results:
+                        st.markdown(result)
+                        
+    except Exception as e:
+        st.error(f"Error generating enhanced analysis: {str(e)}")
 
-class ConstructionDefectAnalyzer:
-    """Main application class"""
+def display_analysis_results(predictions, analysis_time):
+    """Display analysis results with damage details"""
+    st.markdown("### 📊 Analysis Results")
+    st.markdown(f"*Analysis completed in {analysis_time:.2f} seconds*")
     
-    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)
+    detected = False
+    for idx, prob in enumerate(predictions):
+        confidence = float(prob) * 100
+        if confidence > 15:
+            detected = True
+            damage_type = DAMAGE_TYPES[idx]['name']
+            cases = KNOWLEDGE_BASE[damage_type]
             
-            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
-        }
+            with st.expander(f"{damage_type.replace('_', ' ').title()} - {confidence:.1f}%", expanded=True):
+                st.markdown(
+                    f"""
+                    <style>
+                    .stProgress > div > div > div > div {{
+                        background-color: {DAMAGE_TYPES[idx]['color']} !important;
+                    }}
+                    </style>
+                    """,
+                    unsafe_allow_html=True
+                )
+                st.progress(confidence / 100)
+                
+                tabs = st.tabs(["📋 Details", "🔧 Repairs", "⚠️ Actions"])
+                
+                with tabs[0]:
+                    for case in cases:
+                        st.markdown(f"""
+                        - **Severity:** {case['severity']}
+                        - **Description:** {case['description']}
+                        - **Location:** {case['location']}
+                        - **Required Expertise:** {case['required_expertise']}
+                        """)
+                
+                with tabs[1]:
+                    for step in cases[0]['repair_method']:
+                        st.markdown(f"✓ {step}")
+                    st.info(f"**Estimated Cost:** {cases[0]['estimated_cost']}")
+                    st.info(f"**Timeframe:** {cases[0]['timeframe']}")
+                
+                with tabs[2]:
+                    st.warning("**Immediate Actions Required:**")
+                    st.markdown(cases[0]['immediate_action'])
+                    st.success("**Prevention Measures:**")
+                    st.markdown(cases[0]['prevention'])
+            
+            # Display enhanced analysis
+            display_enhanced_analysis(damage_type, confidence)
     
-    @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 + "?"
+    if not detected:
+        st.info("No significant structural damage detected. Regular maintenance recommended.")
 
 def main():
-    st.set_page_config(page_title="Advanced Construction Defect Analyzer", layout="wide")
+    """Main application function"""
+    init_session_state()
+    st.set_page_config(
+        page_title="Structural Damage Analyzer Pro",
+        page_icon="🏗️",
+        layout="wide",
+        initial_sidebar_state="expanded"
+    )
+
+    st.markdown(get_custom_css(), unsafe_allow_html=True)
     
-    analyzer = ConstructionDefectAnalyzer()
+    # Sidebar
+    with st.sidebar:
+        st.markdown("### ⚙️ Settings")
+        st.session_state.dark_mode = st.toggle("Dark Mode", st.session_state.dark_mode)
+        st.markdown("### 📖 Analysis History")
+        if st.session_state.history:
+            for item in st.session_state.history[-5:]:
+                st.markdown(f"- {item}")
     
-    st.title("🏗️ Advanced Construction Defect Analyzer")
+    display_header()
+
+    # Load model and initialize RAG system
+    global MODEL, PROCESSOR
+    if MODEL is None or PROCESSOR is None:
+        with st.spinner("Loading AI model..."):
+            MODEL, PROCESSOR = load_model()
+            if MODEL is None:
+                st.error("Failed to load model. Please refresh the page.")
+                return
     
-    # Multiple image upload
-    uploaded_files = st.file_uploader(
-        "Upload construction images for analysis",
+    init_rag_system()
+
+    # File upload
+    uploaded_file = st.file_uploader(
+        "Drag and drop or click to upload an image",
         type=['jpg', 'jpeg', 'png'],
-        accept_multiple_files=True
+        help="Supported formats: JPG, JPEG, PNG"
     )
-    
-    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])
+
+    if uploaded_file:
+        try:
+            if uploaded_file.size > MAX_FILE_SIZE:
+                st.error("File size too large. Please upload an image smaller than 5MB.")
+                return
                 
-                with col1:
-                    st.image(image, caption=f"Image {idx + 1}", use_column_width=True)
+            image = preprocess_image(uploaded_file)
+            if image is None:
+                return
                 
-                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])
+            validate_image(image)
+            
+            col1, col2 = st.columns([1, 1])
+            
+            with col1:
+                st.image(image, caption="Uploaded Structure", use_container_width=True)
+            
+            with col2:
+                with st.spinner("🔍 Analyzing damage..."):
+                    start_time = time.time()
+                    predictions = analyze_damage(image, MODEL, PROCESSOR)
+                    analysis_time = time.time() - start_time
+                    
+                    if predictions is not None:
+                        display_analysis_results(predictions, analysis_time)
+                        st.session_state.history.append(f"Analyzed image: {uploaded_file.name}")
+
+        except Exception as e:
+            cleanup_memory()
+            st.error(f"Error processing image: {str(e)}")
+            st.info("Please try uploading a different image.")
+
+    # Footer
+    st.markdown("---")
+    st.markdown(
+        """
+        <div style='text-align: center'>
+            <p>🏗️ Structural Damage Analyzer Pro | Built with Streamlit & Transformers</p>
+            <p style='font-size: 0.8rem;'>For professional use only. Always consult with a structural engineer.</p>
+        </div>
+        """,
+        unsafe_allow_html=True
+    )
 
 if __name__ == "__main__":
     main()
-
+```