Update app.py

app.py

@@ -6,308 +6,185 @@ from sentence_transformers import SentenceTransformer
 from PIL import Image
 import torch
 import numpy as np
-from typing import List, Dict, Tuple, Optional, Any
+from typing import List, Dict, Tuple
 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
+from datetime import datetime
+import matplotlib.pyplot as plt
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-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
-        }
-    }
+class RAGSystem:
+    def __init__(self):
+        self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+        self.knowledge_base = self.load_knowledge_base()
+        self.vector_store = self.create_vector_store()
+        self.query_history = []
 
-    @staticmethod
-    def load_config():
-        """Load configuration with fallback to defaults"""
+    def load_knowledge_base(self) -> List[Dict]:
+        """Load and preprocess knowledge base"""
+        # Using a simplified version of your knowledge base
+        kb = {
+            "spalling": [
+                {
+                    "severity": "Critical",
+                    "description": "Severe concrete spalling with exposed reinforcement",
+                    "repair_method": "Remove deteriorated concrete, clean reinforcement",
+                    "estimated_cost": "Very High ($15,000+)",
+                    "immediate_action": "Evacuate area, install support"
+                }
+            ],
+            "structural_cracks": [
+                {
+                    "severity": "High",
+                    "description": "Active structural cracks >5mm width",
+                    "repair_method": "Structural analysis, epoxy injection",
+                    "estimated_cost": "High ($10,000-$20,000)",
+                    "immediate_action": "Install crack monitors"
+                }
+            ]
+        }
+        
+        documents = []
+        for category, items in kb.items():
+            for item in items:
+                doc_text = f"Category: {category}\n"
+                for key, value in item.items():
+                    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"""
+        texts = [doc["text"] for doc in self.knowledge_base]
+        embeddings = self.embedding_model.encode(texts)
+        dimension = embeddings.shape[1]
+        index = faiss.IndexFlatL2(dimension)
+        index.add(np.array(embeddings).astype('float32'))
+        return index
+
+    def get_relevant_context(self, query: str, k: int = 3) -> str:
+        """Retrieve relevant context based on query"""
         try:
-            if os.path.exists('config.json'):
-                with open('config.json', 'r') as f:
-                    config = json.load(f)
-                return {**ConfigManager.DEFAULT_CONFIG, **config}
+            query_embedding = self.embedding_model.encode([query])
+            D, I = self.vector_store.search(np.array(query_embedding).astype('float32'), k)
+            context = "\n\n".join([self.knowledge_base[i]["text"] for i in I[0]])
+            
+            # Log query
+            self.query_history.append({
+                "timestamp": datetime.now().isoformat(),
+                "query": query
+            })
+            
+            return context
         except Exception as e:
-            logger.warning(f"Error loading config: {e}")
-        return ConfigManager.DEFAULT_CONFIG
-
-config = ConfigManager.load_config()
+            logger.error(f"Error retrieving context: {e}")
+            return ""
 
 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.defect_classes = ["spalling", "structural_cracks", "surface_deterioration"]
+        self.model = self._initialize_model()
+        self.processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
         self.history = []
 
-    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 _initialize_model(self):
+        model = ViTForImageClassification.from_pretrained(
+            "google/vit-base-patch16-224",
+            num_labels=len(self.defect_classes),
+            ignore_mismatched_sizes=True
+        )
+        return model.to(self.device)
 
-    def preprocess_image(self, image: Image.Image) -> Dict[str, Any]:
-        """Enhanced image preprocessing with multiple analyses"""
+    def analyze_image(self, image: Image.Image) -> Dict:
+        """Analyze image for defects"""
         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:
-                        logger.error("Image preprocessing failed.")
-                        return None  # Early return if preprocessing failed
+            # Preprocess image
+            inputs = self.processor(images=image, return_tensors="pt").to(self.device)
             
-                    # 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
+            # Get model predictions
+            with torch.no_grad():
+                outputs = self.model(**inputs)
             
-                    # Additional analysis based on image statistics
-                    additional_analysis = self.analyze_image_statistics(proc_data["stats"])
+            # Process results
+            probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
+            defect_probs = {
+                self.defect_classes[i]: float(probabilities[0][i])
+                for i in range(len(self.defect_classes))
+            }
             
-                    # 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 = {}
-                        
-                        # 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"
-                        
-                        # Edge density analysis
-                        if stats["edge_density"] > 0.1:
-                            analysis["crack_likelihood"] = "High crack probability based on edge detection"
-                        
-                        # 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)
+            # Basic image statistics
+            img_array = np.array(image)
+            stats = {
+                "mean_brightness": float(np.mean(img_array)),
+                "image_size": image.size
+            }
             
-            # Apply gamma correction
-            gamma = self.analysis_config["heatmap_intensity"]
-            heatmap = np.power(heatmap, gamma)
+            result = {
+                "defect_probabilities": defect_probs,
+                "image_statistics": stats,
+                "timestamp": datetime.now().isoformat()
+            }
             
-            # Apply colormap
-            heatmap = (heatmap * 255).astype(np.uint8)
-            heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+            self.history.append(result)
+            return result
             
-            return heatmap
         except Exception as e:
-            logger.error(f"Heatmap generation error: {e}")
+            logger.error(f"Image analysis error: {e}")
             return None
-class RAGSystem:
-    """Basic RAG System for storing and retrieving documents."""
-    def __init__(self):
-        self.embedding_model = SentenceTransformer(config["model_settings"]["sentence_transformer"])
-        self.vector_store = faiss.IndexFlatL2(384)  # 384-dim for MiniLM embeddings
-        self.knowledge_base = []
 
-    def add_documents(self, docs: List[str]):
-        """Add documents to the vector store."""
-        embeddings = self.embedding_model.encode(docs)
-        self.vector_store.add(np.array(embeddings).astype('float32'))
-        for doc in docs:
-            self.knowledge_base.append({"text": doc})
-
-    def search(self, query: str, k: int = 3):
-        """Retrieve similar documents for the query."""
-        query_embedding = self.embedding_model.encode([query])
-        D, I = self.vector_store.search(np.array(query_embedding).astype('float32'), k)
-        return [self.knowledge_base[i]["text"] for i in I[0]]
-
-
-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)
+def get_groq_response(query: str, context: str) -> str:
+    """Get response from Groq LLM"""
+    try:
+        client = Groq(api_key=os.getenv("GROQ_API_KEY"))
         
-        # 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)
+        prompt = f"""Based on the following context about construction defects, answer the question.
+        Context: {context}
+        Question: {query}
+        Provide a detailed answer based on the context."""
+
+        response = client.chat.completions.create(
+            messages=[
+                {
+                    "role": "system",
+                    "content": "You are a construction defect analysis expert."
+                },
+                {
+                    "role": "user",
+                    "content": prompt
+                }
+            ],
+            model="llama2-70b-4096",
+            temperature=0.7,
+        )
+        return response.choices[0].message.content
+    except Exception as e:
+        logger.error(f"Groq API error: {e}")
+        return f"Error: Unable to get response from AI model. Please try again later."
 
 def main():
     st.set_page_config(
-        page_title="Enhanced Construction Defect Analyzer",
+        page_title="Construction Defect Analyzer",
         page_icon="🏗️",
         layout="wide"
     )
 
-    st.title("🏗️ Advanced Construction Defect Analysis System")
+    st.title("🏗️ Construction Defect Analyzer")
 
-    # Initialize systems
+    # Initialize systems in session state
     if 'rag_system' not in st.session_state:
-        st.session_state.rag_system = EnhancedRAGSystem()
+        st.session_state.rag_system = RAGSystem()
     if 'image_analyzer' not in st.session_state:
         st.session_state.image_analyzer = ImageAnalyzer()
 
-    # 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:])
-
-    # Main interface
-    col1, col2 = st.columns([2, 3])
+    # Create two columns
+    col1, col2 = st.columns([1, 1])
 
     with col1:
         uploaded_file = st.file_uploader(
@@ -315,61 +192,59 @@ def main():
             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"
-        )
-
-    with col2:
         if uploaded_file:
             image = Image.open(uploaded_file)
+            st.image(image, caption="Uploaded Image", use_column_width=True)
             
-            # Create tabs for different views
-            tabs = st.tabs(["Original", "Analysis", "Details"])
-            
-            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]:
+            with st.spinner("Analyzing image..."):
+                results = st.session_state.image_analyzer.analyze_image(image)
+                
                 if results:
-                    st.json(results["additional_analysis"])
-                    if show_debug:
+                    st.subheader("Detected Defects")
+                    
+                    # Create bar chart
+                    defect_probs = results["defect_probabilities"]
+                    fig, ax = plt.subplots()
+                    defects = list(defect_probs.keys())
+                    probs = list(defect_probs.values())
+                    ax.barh(defects, probs)
+                    ax.set_xlim(0, 1)
+                    ax.set_xlabel("Probability")
+                    st.pyplot(fig)
+                    
+                    # Show image statistics
+                    if st.checkbox("Show Image Details"):
                         st.json(results["image_statistics"])
 
+    with col2:
+        st.subheader("Ask About Defects")
+        user_query = st.text_input(
+            "Enter your question about construction defects:",
+            help="Example: What are the repair methods for severe spalling?"
+        )
+        
         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("AI Response:")
                 st.write(response)
                 
-                if show_debug:
-                    st.subheader("Retrieved Context")
+                if st.checkbox("Show Retrieved Context"):
+                    st.write("Context Used:")
                     st.text(context)
 
+    # Sidebar for history
+    with st.sidebar:
+        st.header("Analysis History")
+        if st.button("Show Recent Analyses"):
+            if st.session_state.image_analyzer.history:
+                for analysis in st.session_state.image_analyzer.history[-5:]:
+                    st.write(f"Analysis from: {analysis['timestamp']}")
+                    st.json(analysis["defect_probabilities"])
+            else:
+                st.write("No analyses yet")
+
 if __name__ == "__main__":
     main()