Update app.py

app.py

@@ -11,6 +11,7 @@ from typing import List, Dict
 from datetime import datetime
 from groq import Groq
 import os
+from functools import lru_cache
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -18,13 +19,34 @@ 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 = []
+        # Load models only when needed
+        self._embedding_model = None
+        self._vector_store = None
+        self._knowledge_base = None
 
-    def load_knowledge_base(self) -> List[Dict]:
+    @property
+    def embedding_model(self):
+        if self._embedding_model is None:
+            self._embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+        return self._embedding_model
+
+    @property
+    def knowledge_base(self):
+        if self._knowledge_base is None:
+            self._knowledge_base = self.load_knowledge_base()
+        return self._knowledge_base
+
+    @property
+    def vector_store(self):
+        if self._vector_store is None:
+            self._vector_store = self.create_vector_store()
+        return self._vector_store
+
+    @staticmethod
+    @lru_cache(maxsize=1)  # Cache the knowledge base
+    def load_knowledge_base() -> List[Dict]:
         """Load and preprocess knowledge base"""
+        # Your existing knowledge base code...
         kb = {
             "spalling": [
                 {
@@ -77,18 +99,13 @@ class RAGSystem:
         index.add(np.array(embeddings).astype('float32'))
         return index
 
+    @lru_cache(maxsize=32)  # Cache recent query results
     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}")
@@ -96,62 +113,73 @@ class RAGSystem:
 
 class ImageAnalyzer:
     def __init__(self):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = "cpu"  # Force CPU usage for better compatibility
         self.defect_classes = ["spalling", "structural_cracks", "surface_deterioration"]
-        
+        self._model = None
+        self._feature_extractor = None
+
+    @property
+    def model(self):
+        if self._model is None:
+            self._model = self._load_model()
+        return self._model
+
+    @property
+    def feature_extractor(self):
+        if self._feature_extractor is None:
+            self._feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
+        return self._feature_extractor
+
+    def _load_model(self):
         try:
-            # Use feature extractor instead of processor
-            self.feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
-            self.model = ViTForImageClassification.from_pretrained(
+            model = ViTForImageClassification.from_pretrained(
                 "google/vit-base-patch16-224",
                 num_labels=len(self.defect_classes),
                 ignore_mismatched_sizes=True
             ).to(self.device)
             
-            # Initialize the model weights for our specific classes
             with torch.no_grad():
-                self.model.classifier = torch.nn.Linear(
-                    in_features=self.model.classifier.in_features,
+                model.classifier = torch.nn.Linear(
+                    in_features=model.classifier.in_features,
                     out_features=len(self.defect_classes)
                 )
-                
+            return model
         except Exception as e:
             logger.error(f"Model initialization error: {e}")
-            self.model = None
-            self.feature_extractor = None
+            return None
 
-    def preprocess_image(self, image):
+    @st.cache_data  # Cache preprocessed images
+    def preprocess_image(self, image_bytes):
         """Preprocess image for model input"""
-        if image.mode != 'RGB':
-            image = image.convert('RGB')
-        
-        # Resize image to expected size
-        width, height = 224, 224
-        image = image.resize((width, height), Image.Resampling.LANCZOS)
-        
-        return image
+        try:
+            image = Image.open(image_bytes)
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+            
+            width, height = 224, 224
+            image = image.resize((width, height), Image.Resampling.LANCZOS)
+            return image
+        except Exception as e:
+            logger.error(f"Image preprocessing error: {e}")
+            return None
 
     def analyze_image(self, image):
         """Analyze image for defects"""
         try:
-            # Preprocess image
-            processed_image = self.preprocess_image(image)
-            
-            # Extract features
+            if self.model is None:
+                raise ValueError("Model not properly initialized")
+
             inputs = self.feature_extractor(
-                images=processed_image,
+                images=image,
                 return_tensors="pt"
             )
             inputs = {k: v.to(self.device) for k, v in inputs.items()}
             
-            # Get predictions
             with torch.no_grad():
                 outputs = self.model(**inputs)
             
-            # Get probabilities
             probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
             
-            # Add confidence threshold
             confidence_threshold = 0.3
             results = {
                 self.defect_classes[i]: float(probs[i]) 
@@ -159,7 +187,6 @@ class ImageAnalyzer:
                 if float(probs[i]) > confidence_threshold
             }
             
-            # If no defects meet threshold, return the highest probability one
             if not results:
                 max_idx = torch.argmax(probs)
                 results = {self.defect_classes[int(max_idx)]: float(probs[max_idx])}
@@ -170,9 +197,13 @@ class ImageAnalyzer:
             logger.error(f"Analysis error: {str(e)}")
             return None
 
+@st.cache_data
 def get_groq_response(query: str, context: str) -> str:
-    """Get response from Groq LLM"""
+    """Get response from Groq LLM with caching"""
     try:
+        if not os.getenv("GROQ_API_KEY"):
+            return "Error: Groq API key not configured"
+
         client = Groq(api_key=os.getenv("GROQ_API_KEY"))
         
         prompt = f"""Based on the following context about construction defects, answer the question.
@@ -197,7 +228,7 @@ def get_groq_response(query: str, context: str) -> str:
         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."
+        return f"Error: Unable to get response from AI model. Please try again."
 
 def main():
     st.set_page_config(
@@ -208,56 +239,61 @@ def main():
     
     st.title("🏗️ Construction Defect Analyzer")
     
-    # Initialize systems with error handling
-    try:
-        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()
-    except Exception as e:
-        st.error(f"Error initializing systems: {str(e)}")
-        return
+    # Initialize systems in session state if not present
+    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()
     
-    # 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"])
+        uploaded_file = st.file_uploader(
+            "Upload a construction image for analysis",
+            type=["jpg", "jpeg", "png"],
+            key="image_uploader"  # Add key for proper state management
+        )
 
         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)
+                # Create a placeholder for the image
+                image_placeholder = st.empty()
                 
-                # Analyze image
-                with st.spinner('Analyzing image...'):
-                    results = st.session_state.analyzer.analyze_image(image)
-                    
-                    if results:
-                        st.success('Analysis complete!')
+                # Process image with progress indicator
+                with st.spinner('Processing image...'):
+                    processed_image = st.session_state.analyzer.preprocess_image(uploaded_file)
+                    if processed_image:
+                        image_placeholder.image(processed_image, caption='Uploaded Image', use_column_width=True)
                         
-                        # Display results
-                        st.subheader("Detected Defects")
+                        # Analyze image with progress bar
+                        progress_bar = st.progress(0)
+                        with st.spinner('Analyzing defects...'):
+                            results = st.session_state.analyzer.analyze_image(processed_image)
+                            progress_bar.progress(100)
                         
-                        # 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}")
+                        if results:
+                            st.success('Analysis complete!')
+                            
+                            # Display results
+                            st.subheader("Detected Defects")
+                            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)
+                            
+                            most_likely_defect = max(results.items(), key=lambda x: x[1])[0]
+                            st.info(f"Most likely defect: {most_likely_defect}")
+                        else:
+                            st.warning("No defects detected or analysis failed. Please try another image.")
                     else:
-                        st.error("Analysis failed. Please try again.")
-                    
+                        st.error("Failed to process image. Please try another one.")
+                        
             except Exception as e:
-                st.error(f"Error: {str(e)}")
+                st.error(f"Error processing image: {str(e)}")
                 logger.error(f"Process error: {e}")
     
     with col2:
@@ -272,18 +308,21 @@ def main():
                 # 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)
+                if context:
+                    # Get response from Groq
+                    response = get_groq_response(user_query, context)
+                    
+                    if not response.startswith("Error"):
+                        st.write("Answer:")
+                        st.markdown(response)
+                    else:
+                        st.error(response)
+                    
+                    with st.expander("View retrieved information"):
+                        st.text(context)
+                else:
+                    st.error("Could not find relevant information. Please try rephrasing your question.")
 
-    # Sidebar for information and settings
     with st.sidebar:
         st.header("About")
         st.write("""
@@ -304,9 +343,9 @@ def main():
         
         # Add settings section
         st.subheader("Settings")
-        if st.button("Clear Session"):
-            st.session_state.clear()
-            st.success("Session cleared!")
+        if st.button("Clear Cache"):
+            st.cache_data.clear()
+            st.success("Cache cleared!")
 
 if __name__ == "__main__":
     main()