Update app.py

app.py

@@ -6,416 +6,359 @@ from sentence_transformers import SentenceTransformer
 from PIL import Image
 import torch
 import numpy as np
-from typing import List, Dict
+from typing import List, Dict, Tuple, Optional, Any
 import faiss
 import json
+import torchvision.transforms.functional as TF
+from torchvision import transforms
+import cv2
+import pandas as pd
+from datetime import datetime
+import logging
 
-# Initialize sentence transformer for embeddings
-@st.cache_resource
-def init_embedding_model():
-    return SentenceTransformer('all-MiniLM-L6-v2')
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Initialize Groq client
-@st.cache_resource
-def init_groq_client():
-    return Groq(api_key=os.environ.get("GROQ_API_KEY"))
-class RAGSystem:
-    def __init__(self):
-        self.embedding_model = init_embedding_model()
-        self.knowledge_base = self.load_knowledge_base()
-        self.vector_store = self.create_vector_store()
-        
-    def load_knowledge_base(self) -> List[Dict]:
-        """Load and preprocess knowledge base into a list of documents"""
-        # Your existing knowledge base dictionary
-        kb = {
-            "spalling": [
-                {
-                    "severity": "Critical",
-                    "description": "Severe concrete spalling with exposed reinforcement and section loss",
-                    "repair_method": [
-                        "Install temporary support",
-                        "Remove deteriorated concrete",
-                        "Clean and treat reinforcement",
-                        "Apply corrosion inhibitor",
-                        "Apply bonding agent",
-                        "High-strength repair mortar",
-                        "Surface treatment and waterproofing"
-                    ],
-                    "estimated_cost": "Very High ($15,000+)",
-                    "timeframe": "3-4 weeks",
-                    "location": "Primary structural elements",
-                    "required_expertise": "Structural Engineer + Specialist Contractor",
-                    "immediate_action": "Evacuate area, install temporary support, prevent access",
-                    "prevention": "Regular inspections, waterproofing, chloride protection",
-                    "testing_required": ["Core testing", "Reinforcement scanning", "Chloride testing"],
-                    "common_causes": [
-                        "Reinforcement corrosion",
-                        "Freeze-thaw cycles",
-                        "Poor concrete cover",
-                        "Chemical attack"
-                    ],
-                    "safety_considerations": [
-                        "Risk of structural failure",
-                        "Falling concrete hazard",
-                        "Worker safety during repairs"
-                    ]
-                },
-                {
-                    "severity": "Moderate",
-                    "description": "Surface spalling without exposed reinforcement",
-                    "repair_method": [
-                        "Remove loose concrete",
-                        "Surface preparation",
-                        "Apply repair mortar",
-                        "Surface treatment"
-                    ],
-                    "estimated_cost": "Medium ($5,000-$10,000)",
-                    "timeframe": "1-2 weeks",
-                    "location": "Non-structural elements",
-                    "required_expertise": "Concrete Repair Specialist",
-                    "immediate_action": "Remove loose material, protect from water ingress",
-                    "prevention": "Surface sealers, proper drainage",
-                    "testing_required": ["Surface adhesion testing", "Moisture testing"],
-                    "common_causes": [
-                        "Surface carbonation",
-                        "Impact damage",
-                        "Poor curing"
-                    ],
-                    "safety_considerations": [
-                        "Minor falling debris risk",
-                        "Dust control during repairs"
-                    ]
-                }
-            ],
-            "reinforcement_corrosion": [
-                {
-                    "severity": "Critical",
-                    "description": "Severe corrosion with >30% section loss",
-                    "repair_method": [
-                        "Structural support installation",
-                        "Concrete removal around reinforcement",
-                        "Reinforcement replacement",
-                        "Corrosion protection application",
-                        "High-strength concrete repair",
-                        "Cathodic protection installation"
-                    ],
-                    "estimated_cost": "Critical ($20,000+)",
-                    "timeframe": "4-6 weeks",
-                    "location": "Load-bearing elements",
-                    "required_expertise": "Senior Structural Engineer",
-                    "immediate_action": "Immediate evacuation, emergency shoring",
-                    "prevention": "Waterproofing, cathodic protection",
-                    "testing_required": [
-                        "Half-cell potential survey",
-                        "Concrete resistivity testing",
-                        "Chloride analysis",
-                        "Carbonation testing"
-                    ],
-                    "common_causes": [
-                        "Chloride contamination",
-                        "Carbonation",
-                        "Stray electrical currents",
-                        "Poor concrete quality"
-                    ],
-                    "safety_considerations": [
-                        "Structural collapse risk",
-                        "Electrical hazards during testing",
-                        "Confined space entry"
-                    ]
-                }
-            ],
-            "structural_cracks": [
-                {
-                    "severity": "High",
-                    "description": "Active structural cracks >5mm width",
-                    "repair_method": [
-                        "Structural analysis",
-                        "Crack monitoring",
-                        "Epoxy injection",
-                        "Carbon fiber reinforcement",
-                        "Load path modification"
-                    ],
-                    "estimated_cost": "High ($10,000-$20,000)",
-                    "timeframe": "2-4 weeks",
-                    "location": "Primary structural elements",
-                    "required_expertise": "Structural Engineer",
-                    "immediate_action": "Install crack monitors, restrict loading",
-                    "prevention": "Proper design, joint maintenance",
-                    "testing_required": [
-                        "Crack movement monitoring",
-                        "Load testing",
-                        "Concrete strength testing"
-                    ],
-                    "common_causes": [
-                        "Overloading",
-                        "Foundation settlement",
-                        "Thermal movements",
-                        "Design deficiencies"
-                    ],
-                    "safety_considerations": [
-                        "Structural stability",
-                        "Water infiltration",
-                        "Working at height"
-                    ]
-                }
-            ],
-            "water_damage": [
-                {
-                    "severity": "Medium",
-                    "description": "Active water infiltration with deterioration",
-                    "repair_method": [
-                        "Water source identification",
-                        "Drainage improvement",
-                        "Waterproofing membrane installation",
-                        "Joint sealing",
-                        "Surface treatment"
-                    ],
-                    "estimated_cost": "Medium ($5,000-$15,000)",
-                    "timeframe": "1-3 weeks",
-                    "location": "Various locations",
-                    "required_expertise": "Waterproofing Specialist",
-                    "immediate_action": "Water diversion, dehumidification",
-                    "prevention": "Regular maintenance, proper drainage",
-                    "testing_required": [
-                        "Moisture mapping",
-                        "Drainage assessment",
-                        "Permeability testing"
-                    ],
-                    "common_causes": [
-                        "Failed waterproofing",
-                        "Poor drainage",
-                        "Joint failure",
-                        "Condensation"
-                    ],
-                    "safety_considerations": [
-                        "Slip hazards",
-                        "Electrical safety",
-                        "Mold growth"
-                    ]
-                }
-            ],
-            "surface_deterioration": [
-                {
-                    "severity": "Low",
-                    "description": "Surface scaling and deterioration",
-                    "repair_method": [
-                        "Surface cleaning",
-                        "Repair material application",
-                        "Surface treatment",
-                        "Protective coating"
-                    ],
-                    "estimated_cost": "Low ($2,000-$5,000)",
-                    "timeframe": "3-5 days",
-                    "location": "Exposed surfaces",
-                    "required_expertise": "Concrete Repair Technician",
-                    "immediate_action": "Clean and protect surface",
-                    "prevention": "Regular maintenance, surface protection",
-                    "testing_required": [
-                        "Surface strength testing",
-                        "Coating adhesion tests"
-                    ],
-                    "common_causes": [
-                        "Freeze-thaw damage",
-                        "Chemical exposure",
-                        "Poor finishing",
-                        "Abrasion"
-                    ],
-                    "safety_considerations": [
-                        "Dust control",
-                        "Chemical handling",
-                        "PPE requirements"
-                    ]
-                }
-            ],
-            "alkali_silica_reaction": [
-                {
-                    "severity": "High",
-                    "description": "Concrete expansion and map cracking due to ASR",
-                    "repair_method": [
-                        "Expansion monitoring",
-                        "Moisture control",
-                        "Crack sealing",
-                        "Surface treatment",
-                        "Structural strengthening"
-                    ],
-                    "estimated_cost": "High ($15,000-$25,000)",
-                    "timeframe": "3-5 weeks",
-                    "location": "Concrete elements",
-                    "required_expertise": "Materials Engineer + Structural Engineer",
-                    "immediate_action": "Monitor expansion, control moisture",
-                    "prevention": "Proper aggregate selection, pozzolans",
-                    "testing_required": [
-                        "Petrographic analysis",
-                        "Expansion testing",
-                        "Humidity monitoring"
-                    ],
-                    "common_causes": [
-                        "Reactive aggregates",
-                        "High alkali cement",
-                        "Moisture presence",
-                        "Temperature cycles"
-                    ],
-                    "safety_considerations": [
-                        "Progressive deterioration",
-                        "Structural integrity",
-                        "Long-term monitoring"
-                    ]
-                }
-            ]
+class ConfigManager:
+    """Manages configuration settings for the application"""
+    DEFAULT_CONFIG = {
+        "model_settings": {
+            "vit_model": "google/vit-base-patch16-224",
+            "sentence_transformer": "all-MiniLM-L6-v2",
+            "groq_model": "llama-3.3-70b-versatile"
+        },
+        "analysis_settings": {
+            "confidence_threshold": 0.5,
+            "max_defects": 3,
+            "heatmap_intensity": 0.7
+        },
+        "rag_settings": {
+            "num_relevant_docs": 3,
+            "similarity_threshold": 0.75
         }
-        
-        # Convert nested knowledge base into flat documents
-        documents = []
-        for category, items in kb.items():
-            for item in items:
-                # Create a text representation of the document
-                doc_text = f"Category: {category}\n"
-                for key, value in item.items():
-                    if isinstance(value, list):
-                        doc_text += f"{key}: {', '.join(value)}\n"
-                    else:
-                        doc_text += f"{key}: {value}\n"
-                documents.append({
-                    "text": doc_text,
-                    "metadata": {"category": category}
-                })
-        
-        return documents
+    }
 
-    def create_vector_store(self):
-        """Create FAISS vector store from knowledge base"""
-        # Generate embeddings for all documents
-        texts = [doc["text"] for doc in self.knowledge_base]
-        embeddings = self.embedding_model.encode(texts)
-        
-        # Initialize FAISS index
-        dimension = embeddings.shape[1]
-        index = faiss.IndexFlatL2(dimension)
-        index.add(np.array(embeddings).astype('float32'))
-        
-        return index
-    
-    def create_vector_store(self):
-        """Create FAISS vector store from knowledge base"""
-        # Generate embeddings for all documents
-        texts = [doc["text"] for doc in self.knowledge_base]
-        embeddings = self.embedding_model.encode(texts)
+    @staticmethod
+    def load_config():
+        """Load configuration with fallback to defaults"""
+        try:
+            if os.path.exists('config.json'):
+                with open('config.json', 'r') as f:
+                    config = json.load(f)
+                return {**ConfigManager.DEFAULT_CONFIG, **config}
+        except Exception as e:
+            logger.warning(f"Error loading config: {e}")
+        return ConfigManager.DEFAULT_CONFIG
+
+config = ConfigManager.load_config()
+
+class ImageAnalyzer:
+    def __init__(self):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.config = config["model_settings"]
+        self.analysis_config = config["analysis_settings"]
+        self.defect_classes = [
+            "spalling", "reinforcement_corrosion", "structural_cracks",
+            "water_damage", "surface_deterioration", "alkali_silica_reaction",
+            "concrete_delamination", "honeycomb", "scaling",
+            "efflorescence", "joint_deterioration", "carbonation"
+        ]
+        self.initialize_models()
+        self.history = []
+
+    @st.cache_resource
+    def initialize_models(self):
+        """Initialize all required models"""
+        try:
+            # Initialize ViT model
+            self.model = ViTForImageClassification.from_pretrained(
+                self.config["vit_model"],
+                num_labels=len(self.defect_classes),
+                ignore_mismatched_sizes=True
+            ).to(self.device)
+            
+            # Initialize image processor
+            self.processor = ViTImageProcessor.from_pretrained(self.config["vit_model"])
+            
+            # Initialize transformations pipeline
+            self.transforms = self._setup_transforms()
+            
+            return True
+        except Exception as e:
+            logger.error(f"Model initialization error: {e}")
+            return False
+
+    def _setup_transforms(self):
+        """Setup image transformation pipeline"""
+        return transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], 
+                              std=[0.229, 0.224, 0.225]),
+            transforms.RandomAdjustSharpness(2),
+            transforms.ColorJitter(brightness=0.2, contrast=0.2)
+        ])
+
+    def preprocess_image(self, image: Image.Image) -> Dict[str, Any]:
+        """Enhanced image preprocessing with multiple analyses"""
+        try:
+            # Convert to RGB if necessary
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+
+            # Basic image statistics
+            img_array = np.array(image)
+            stats = {
+                "mean_brightness": np.mean(img_array),
+                "std_brightness": np.std(img_array),
+                "size": image.size,
+                "aspect_ratio": image.size[0] / image.size[1]
+            }
+
+            # Edge detection for crack analysis
+            gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
+            edges = cv2.Canny(gray, 100, 200)
+            stats["edge_density"] = np.mean(edges > 0)
+
+            # Color analysis for rust detection
+            hsv = cv2.cvtColor(img_array, cv2.COLOR_RGB2HSV)
+            rust_mask = cv2.inRange(hsv, np.array([0, 50, 50]), np.array([30, 255, 255]))
+            stats["rust_percentage"] = np.mean(rust_mask > 0)
+
+            # Transform for model
+            model_input = self.transforms(image).unsqueeze(0).to(self.device)
+
+            return {
+                "model_input": model_input,
+                "stats": stats,
+                "edges": edges,
+                "rust_mask": rust_mask
+            }
+        except Exception as e:
+            logger.error(f"Preprocessing error: {e}")
+            return None
+
+    def detect_defects(self, image: Image.Image) -> Dict[str, Any]:
+        """Enhanced defect detection with multiple analysis methods"""
+        try:
+            # Preprocess image
+            proc_data = self.preprocess_image(image)
+            if proc_data is None:
+                return None
+
+            # Model prediction
+            with torch.no_grad():
+                outputs = self.model(proc_data["model_input"])
+                
+            # Get probabilities
+            probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
+            
+            # Convert to dictionary
+            defect_probs = {
+                self.defect_classes[i]: float(probabilities[0][i])
+                for i in range(len(self.defect_classes))
+            }
+
+            # Generate attention heatmap
+            attention_weights = outputs.attentions[-1].mean(dim=1)[0] if hasattr(outputs, 'attentions') else None
+            heatmap = self.generate_heatmap(attention_weights, image.size) if attention_weights is not None else None
+
+            # Additional analysis based on image statistics
+            additional_analysis = self.analyze_image_statistics(proc_data["stats"])
+
+            # Combine all results
+            result = {
+                "defect_probabilities": defect_probs,
+                "heatmap": heatmap,
+                "image_statistics": proc_data["stats"],
+                "additional_analysis": additional_analysis,
+                "edge_detection": proc_data["edges"],
+                "rust_detection": proc_data["rust_mask"],
+                "timestamp": datetime.now().isoformat()
+            }
+
+            # Save to history
+            self.history.append(result)
+
+            return result
+        except Exception as e:
+            logger.error(f"Defect detection error: {e}")
+            return None
+
+    def analyze_image_statistics(self, stats: Dict) -> Dict[str, Any]:
+        """Analyze image statistics for additional insights"""
+        analysis = {}
         
-        # Initialize FAISS index
-        dimension = embeddings.shape[1]
-        index = faiss.IndexFlatL2(dimension)
-        index.add(np.array(embeddings).astype('float32'))
+        # Brightness analysis
+        if stats["mean_brightness"] < 50:
+            analysis["lighting_condition"] = "Poor lighting - may affect accuracy"
+        elif stats["mean_brightness"] > 200:
+            analysis["lighting_condition"] = "Overexposed - may affect accuracy"
         
-        return index
+        # Edge density analysis
+        if stats["edge_density"] > 0.1:
+            analysis["crack_likelihood"] = "High crack probability based on edge detection"
         
-    def get_relevant_context(self, query: str, k: int = 3) -> str:
-        """Retrieve relevant context based on query"""
+        # Rust analysis
+        if stats["rust_percentage"] > 0.05:
+            analysis["corrosion_indicator"] = "Possible corrosion detected"
+
+        return analysis
+
+    def generate_heatmap(self, attention_weights: torch.Tensor, image_size: Tuple[int, int]) -> np.ndarray:
+        """Generate enhanced attention heatmap"""
+        try:
+            if attention_weights is None:
+                return None
+
+            # Process attention weights
+            heatmap = attention_weights.cpu().numpy()
+            heatmap = cv2.resize(heatmap, image_size)
+            
+            # Enhanced normalization
+            heatmap = np.maximum(heatmap, 0)
+            heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min() + 1e-8)
+            
+            # Apply gamma correction
+            gamma = self.analysis_config["heatmap_intensity"]
+            heatmap = np.power(heatmap, gamma)
+            
+            # Apply colormap
+            heatmap = (heatmap * 255).astype(np.uint8)
+            heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+            
+            return heatmap
+        except Exception as e:
+            logger.error(f"Heatmap generation error: {e}")
+            return None
+
+class EnhancedRAGSystem(RAGSystem):
+    """Enhanced RAG system with additional features"""
+    def __init__(self):
+        super().__init__()
+        self.config = config["rag_settings"]
+        self.query_history = []
+
+    def get_relevant_context(self, query: str, k: int = None) -> str:
+        """Enhanced context retrieval with debugging info"""
+        if k is None:
+            k = self.config["num_relevant_docs"]
+
+        # Log query
+        self.query_history.append({
+            "timestamp": datetime.now().isoformat(),
+            "query": query
+        })
+
         # Generate query embedding
         query_embedding = self.embedding_model.encode([query])
         
         # Search for similar documents
         D, I = self.vector_store.search(np.array(query_embedding).astype('float32'), k)
         
-        # Combine relevant documents into context
-        context = "\n\n".join([self.knowledge_base[i]["text"] for i in I[0]])
-        return context
-
-    def get_groq_response(query: str, context: str) -> str:
-    """Get response from Groq LLM"""
-    client = init_groq_client()
-    try:
-        prompt = f"""Based on the following context about construction defects, please answer the question.
-Context:
-{context}
-Question: {query}
-Please provide a detailed and specific answer based on the given context."""
-
-        response = client.chat.completions.create(
-            messages=[
-                {
-                    "role": "system",
-                    "content": "You are a construction defect analysis expert. Provide detailed, accurate answers based on the given context."
-                },
-                {
-                    "role": "user",
-                    "content": prompt
-                }
-            ],
-            model="llama-3.3-70b-versatile",
-        )
-        return response.choices[0].message.content
-    except Exception as e:
-        return f"Error: {str(e)}"
+        # Filter by similarity threshold
+        relevant_docs = [
+            self.knowledge_base[i]["text"] 
+            for i, dist in zip(I[0], D[0]) 
+            if dist < self.config["similarity_threshold"]
+        ]
+        
+        return "\n\n".join(relevant_docs)
 
-ef main():
+def main():
     st.set_page_config(
-        page_title="Construction Defect RAG Analyzer",
+        page_title="Enhanced Construction Defect Analyzer",
         page_icon="🏗️",
         layout="wide"
     )
 
-    st.title("🏗️ Construction Defect RAG Analyzer")
+    st.title("🏗️ Advanced Construction Defect Analysis System")
 
-    # Initialize RAG system
+    # Initialize systems
     if 'rag_system' not in st.session_state:
-        st.session_state.rag_system = RAGSystem()
+        st.session_state.rag_system = EnhancedRAGSystem()
+    if 'image_analyzer' not in st.session_state:
+        st.session_state.image_analyzer = ImageAnalyzer()
 
-    # File upload for image analysis
-    uploaded_file = st.file_uploader(
-        "Upload a construction image",
-        type=['jpg', 'jpeg', 'png']
-    )
+    # Sidebar for settings and history
+    with st.sidebar:
+        st.header("Settings & History")
+        show_debug = st.checkbox("Show Debug Information")
+        confidence_threshold = st.slider(
+            "Confidence Threshold",
+            min_value=0.0,
+            max_value=1.0,
+            value=config["analysis_settings"]["confidence_threshold"]
+        )
+        
+        if st.button("View Analysis History"):
+            st.write("Recent Analyses:", st.session_state.image_analyzer.history[-5:])
 
-    # Query input
-    user_query = st.text_input("Ask a question about construction defects:")
+    # Main interface
+    col1, col2 = st.columns([2, 3])
+
+    with col1:
+        uploaded_file = st.file_uploader(
+            "Upload a construction image",
+            type=['jpg', 'jpeg', 'png']
+        )
+        
+        user_query = st.text_input(
+            "Ask a question about construction defects:",
+            help="Enter your question about specific defects or general construction issues"
+        )
 
-    if user_query:
-        with st.spinner("Processing query..."):
-            # Get relevant context using RAG
-            context = st.session_state.rag_system.get_relevant_context(user_query)
+    with col2:
+        if uploaded_file:
+            image = Image.open(uploaded_file)
             
-            # Debug view of retrieved context
-            if st.checkbox("Show retrieved context"):
-                st.subheader("Retrieved Context")
-                st.text(context)
+            # Create tabs for different views
+            tabs = st.tabs(["Original", "Analysis", "Details"])
             
-            # Get response from Groq
-            st.subheader("AI Assistant Response")
-            response = get_groq_response(user_query, context)
-            st.write(response)
-
-    if uploaded_file:
-        image = Image.open(uploaded_file)
-        st.image(image, caption="Uploaded Image")
-        
-        # Combine image analysis with RAG
-        col1, col2 = st.columns(2)
-        
-        with col1:
-            st.subheader("Image Analysis")
-            # Image analysis placeholder
-            st.info("Image analysis results would appear here")
-
-        with col2:
-            st.subheader("AI Assistant Response")
-            if user_query:  # Only show response if there's a query
-                # Get relevant context from knowledge base
+            with tabs[0]:
+                st.image(image, caption="Uploaded Image")
+            
+            with tabs[1]:
+                with st.spinner("Analyzing image..."):
+                    results = st.session_state.image_analyzer.detect_defects(image)
+                    
+                    if results:
+                        # Show defect probabilities
+                        defect_probs = results["defect_probabilities"]
+                        significant_defects = {
+                            k: v for k, v in defect_probs.items() 
+                            if v > confidence_threshold
+                        }
+                        
+                        if significant_defects:
+                            st.subheader("Detected Defects")
+                            fig = plt.figure(figsize=(10, 6))
+                            plt.barh(list(significant_defects.keys()),
+                                   list(significant_defects.values()))
+                            st.pyplot(fig)
+                        
+                        # Show heatmap
+                        if results["heatmap"] is not None:
+                            st.image(results["heatmap"], caption="Defect Attention Map")
+
+            with tabs[2]:
+                if results:
+                    st.json(results["additional_analysis"])
+                    if show_debug:
+                        st.json(results["image_statistics"])
+
+        if user_query:
+            with st.spinner("Processing query..."):
                 context = st.session_state.rag_system.get_relevant_context(user_query)
+                response = get_groq_response(user_query, context)
+                
+                st.subheader("AI Assistant Response")
+                st.write(response)
                 
-                # Get response from Groq
-                with st.spinner("Getting AI response..."):
-                    response = get_groq_response(user_query, context)
-                    st.write(response)
-
-        # Display knowledge base sections
-        if st.checkbox("Show Knowledge Base"):
-            st.subheader("Available Knowledge Base")
-            kb_data = st.session_state.rag_system.knowledge_base
-            for doc in kb_data:
-                category = doc["metadata"]["category"]
-                with st.expander(category.title()):
-                    st.text(doc["text"])
+                if show_debug:
+                    st.subheader("Retrieved Context")
+                    st.text(context)
 
 if __name__ == "__main__":
     main()