Update app.py

app.py

@@ -18,6 +18,32 @@ import matplotlib.pyplot as plt
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Ensure CUDA is available or use CPU
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+@st.cache_resource
+def load_vit_model():
+    """Load and cache the ViT model"""
+    try:
+        model = ViTForImageClassification.from_pretrained(
+            "google/vit-base-patch16-224",
+            num_labels=3,  # Number of defect classes
+            ignore_mismatched_sizes=True
+        )
+        return model.to(DEVICE)
+    except Exception as e:
+        logger.error(f"Error loading ViT model: {e}")
+        return None
+
+@st.cache_resource
+def load_vit_processor():
+    """Load and cache the ViT processor"""
+    try:
+        return ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
+    except Exception as e:
+        logger.error(f"Error loading ViT processor: {e}")
+        return None
+
 class RAGSystem:
     def __init__(self):
         self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
@@ -27,7 +53,6 @@ class RAGSystem:
 
     def load_knowledge_base(self) -> List[Dict]:
         """Load and preprocess knowledge base"""
-        # Using a simplified version of your knowledge base
         kb = {
             "spalling": [
                 {
@@ -75,7 +100,6 @@ class RAGSystem:
             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
@@ -88,26 +112,37 @@ class RAGSystem:
 
 class ImageAnalyzer:
     def __init__(self):
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = DEVICE
         self.defect_classes = ["spalling", "structural_cracks", "surface_deterioration"]
-        self.model = self._initialize_model()
-        self.processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
+        self.model = load_vit_model()
+        self.processor = load_vit_processor()
         self.history = []
 
-    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) -> torch.Tensor:
+        """Preprocess image for model input"""
+        try:
+            # Ensure image is RGB
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+            
+            # Process image using ViT processor
+            inputs = self.processor(images=image, return_tensors="pt")
+            return inputs.to(self.device)
+        except Exception as e:
+            logger.error(f"Image preprocessing error: {e}")
+            return None
 
     def analyze_image(self, image: Image.Image) -> Dict:
         """Analyze image for defects"""
         try:
+            if self.model is None or self.processor is None:
+                raise ValueError("Model or processor not properly initialized")
+
             # Preprocess image
-            inputs = self.processor(images=image, return_tensors="pt").to(self.device)
-            
+            inputs = self.preprocess_image(image)
+            if inputs is None:
+                raise ValueError("Image preprocessing failed")
+
             # Get model predictions
             with torch.no_grad():
                 outputs = self.model(**inputs)
@@ -123,7 +158,7 @@ class ImageAnalyzer:
             img_array = np.array(image)
             stats = {
                 "mean_brightness": float(np.mean(img_array)),
-                "image_size": image.size
+                "image_size": f"{image.size[0]}x{image.size[1]}"
             }
             
             result = {
@@ -182,29 +217,51 @@ def main():
         st.session_state.rag_system = RAGSystem()
     if 'image_analyzer' not in st.session_state:
         st.session_state.image_analyzer = ImageAnalyzer()
+    if 'processed_images' not in st.session_state:
+        st.session_state.processed_images = {}
 
     # Create two columns
     col1, col2 = st.columns([1, 1])
 
     with col1:
+        st.header("Image Analysis")
         uploaded_file = st.file_uploader(
             "Upload a construction image",
-            type=['jpg', 'jpeg', 'png']
+            type=['jpg', 'jpeg', 'png'],
+            key="image_uploader"
         )
         
-        if uploaded_file:
-            image = Image.open(uploaded_file)
-            st.image(image, caption="Uploaded Image", use_column_width=True)
-            
-            with st.spinner("Analyzing image..."):
+        if uploaded_file is not None:
+            try:
+                # Display upload progress
+                progress_bar = st.progress(0)
+                status_text = st.empty()
+                
+                # Update progress
+                status_text.text("Loading image...")
+                progress_bar.progress(25)
+                
+                # Load and display image
+                image = Image.open(uploaded_file)
+                st.image(image, caption="Uploaded Image", use_column_width=True)
+                
+                # Update progress
+                status_text.text("Analyzing image...")
+                progress_bar.progress(50)
+                
+                # Analyze image
                 results = st.session_state.image_analyzer.analyze_image(image)
+                progress_bar.progress(75)
                 
                 if results:
+                    status_text.text("Analysis complete!")
+                    progress_bar.progress(100)
+                    
                     st.subheader("Detected Defects")
                     
                     # Create bar chart
                     defect_probs = results["defect_probabilities"]
-                    fig, ax = plt.subplots()
+                    fig, ax = plt.subplots(figsize=(8, 4))
                     defects = list(defect_probs.keys())
                     probs = list(defect_probs.values())
                     ax.barh(defects, probs)
@@ -213,11 +270,18 @@ def main():
                     st.pyplot(fig)
                     
                     # Show image statistics
-                    if st.checkbox("Show Image Details"):
+                    with st.expander("Image Details"):
                         st.json(results["image_statistics"])
+                else:
+                    status_text.text("Analysis failed. Please try again.")
+                    progress_bar.empty()
+                    
+            except Exception as e:
+                st.error(f"Error processing image: {str(e)}")
+                logger.error(f"Image processing error: {e}")
 
     with col2:
-        st.subheader("Ask About Defects")
+        st.header("Ask About Defects")
         user_query = st.text_input(
             "Enter your question about construction defects:",
             help="Example: What are the repair methods for severe spalling?"
@@ -228,11 +292,10 @@ def main():
                 context = st.session_state.rag_system.get_relevant_context(user_query)
                 response = get_groq_response(user_query, context)
                 
-                st.write("AI Response:")
+                st.subheader("AI Response")
                 st.write(response)
                 
-                if st.checkbox("Show Retrieved Context"):
-                    st.write("Context Used:")
+                with st.expander("View Context Used"):
                     st.text(context)
 
     # Sidebar for history