Update app.py

app.py

@@ -3,19 +3,102 @@ import torch
 from PIL import Image
 import numpy as np
 from transformers import ViTForImageClassification, ViTImageProcessor
+from sentence_transformers import SentenceTransformer
 import matplotlib.pyplot as plt
 import logging
+import faiss
+from typing import List, Dict
+from datetime import datetime
+from groq import Groq
+import os
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+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 = []
+
+    def load_knowledge_base(self) -> List[Dict]:
+        """Load and preprocess 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",
+                    "prevention": "Regular inspections, waterproofing"
+                }
+            ],
+            "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",
+                    "prevention": "Regular monitoring, load management"
+                }
+            ],
+            "surface_deterioration": [
+                {
+                    "severity": "Medium",
+                    "description": "Surface scaling and deterioration",
+                    "repair_method": "Surface preparation, patch repair",
+                    "estimated_cost": "Medium ($5,000-$10,000)",
+                    "immediate_action": "Document extent, plan repairs",
+                    "prevention": "Surface sealers, proper drainage"
+                }
+            ]
+        }
+        
+        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 = 2) -> str:
+        """Retrieve relevant context based on query"""
+        try:
+            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]])
+            
+            self.query_history.append({
+                "timestamp": datetime.now().isoformat(),
+                "query": query
+            })
+            
+            return context
+        except Exception as e:
+            logger.error(f"Error retrieving context: {e}")
+            return ""
+
 class ImageAnalyzer:
     def __init__(self):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.defect_classes = ["spalling", "structural_cracks", "surface_deterioration"]
         
-        # Load model and processor
         try:
             self.model = ViTForImageClassification.from_pretrained(
                 "google/vit-base-patch16-224",
@@ -49,50 +132,133 @@ class ImageAnalyzer:
             logger.error(f"Analysis error: {e}")
             return None
 
+def get_groq_response(query: str, context: str) -> str:
+    """Get response from Groq LLM"""
+    try:
+        client = Groq(api_key=os.getenv("GROQ_API_KEY"))
+        
+        prompt = f"""Based on the following context about construction defects, answer the question.
+        Context: {context}
+        Question: {query}
+        Provide a detailed answer based on the given 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 check your API key and try again."
+
 def main():
-    st.title("Construction Defect Analyzer")
+    st.set_page_config(
+        page_title="Construction Defect Analyzer",
+        page_icon="🏗️",
+        layout="wide"
+    )
+    
+    st.title("🏗️ Construction Defect Analyzer")
     
-    # Initialize analyzer
+    # Initialize systems
     if 'analyzer' not in st.session_state:
         st.session_state.analyzer = ImageAnalyzer()
+    if 'rag_system' not in st.session_state:
+        st.session_state.rag_system = RAGSystem()
     
-    # File uploader
-    st.write("Upload a construction image for analysis")
-    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+    # Create two columns
+    col1, col2 = st.columns([1, 1])
+    
+    with col1:
+        st.subheader("Image Analysis")
+        uploaded_file = st.file_uploader("Upload a construction image for analysis", type=["jpg", "jpeg", "png"])
 
-    if uploaded_file is not None:
-        try:
-            # Display confirmation
-            st.write("Image received. Processing...")
-            
-            # Read and display image
-            image = Image.open(uploaded_file)
-            st.image(image, caption='Uploaded Image', use_column_width=True)
-            
-            # Analyze image
-            with st.spinner('Analyzing image...'):
-                results = st.session_state.analyzer.analyze_image(image)
+        if uploaded_file is not None:
+            try:
+                # Read and display image
+                image = Image.open(uploaded_file)
+                st.image(image, caption='Uploaded Image', use_column_width=True)
                 
-                if results:
-                    st.success('Analysis complete!')
+                # Analyze image
+                with st.spinner('Analyzing image...'):
+                    results = st.session_state.analyzer.analyze_image(image)
                     
-                    # Display results
-                    st.subheader("Defect Probabilities")
+                    if results:
+                        st.success('Analysis complete!')
+                        
+                        # Display results
+                        st.subheader("Detected Defects")
+                        
+                        # Create bar chart
+                        fig, ax = plt.subplots(figsize=(8, 4))
+                        defects = list(results.keys())
+                        probs = list(results.values())
+                        ax.barh(defects, probs)
+                        ax.set_xlim(0, 1)
+                        plt.tight_layout()
+                        st.pyplot(fig)
+                        
+                        # Get most likely defect
+                        most_likely_defect = max(results.items(), key=lambda x: x[1])[0]
+                        st.info(f"Most likely defect: {most_likely_defect}")
+                    else:
+                        st.error("Analysis failed. Please try again.")
                     
-                    # Create bar chart
-                    fig, ax = plt.subplots()
-                    defects = list(results.keys())
-                    probs = list(results.values())
-                    ax.barh(defects, probs)
-                    ax.set_xlim(0, 1)
-                    plt.tight_layout()
-                    st.pyplot(fig)
-                else:
-                    st.error("Analysis failed. Please try again.")
-                    
-        except Exception as e:
-            st.error(f"Error: {str(e)}")
-            logger.error(f"Process error: {e}")
+            except Exception as e:
+                st.error(f"Error: {str(e)}")
+                logger.error(f"Process error: {e}")
+    
+    with col2:
+        st.subheader("Ask About Defects")
+        user_query = st.text_input(
+            "Ask a question about the defects or repairs:",
+            help="Example: What are the repair methods for spalling?"
+        )
+        
+        if user_query:
+            with st.spinner('Getting answer...'):
+                # Get context from RAG system
+                context = st.session_state.rag_system.get_relevant_context(user_query)
+                
+                # Get response from Groq
+                response = get_groq_response(user_query, context)
+                
+                # Display response
+                st.write("Answer:")
+                st.write(response)
+                
+                # Option to view context
+                with st.expander("View retrieved information"):
+                    st.text(context)
+
+    # Sidebar for information
+    with st.sidebar:
+        st.header("About")
+        st.write("""
+        This tool helps analyze construction defects in images and provides 
+        information about repair methods and best practices.
+        
+        Features:
+        - Image analysis for defect detection
+        - Information lookup for repair methods
+        - Expert AI responses to your questions
+        """)
+        
+        # Display API status
+        if os.getenv("GROQ_API_KEY"):
+            st.success("Groq API: Connected")
+        else:
+            st.error("Groq API: Not configured")
 
 if __name__ == "__main__":
     main()