Update app.py

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import torch
 from PIL import Image
 import numpy as np
-from transformers import ViTForImageClassification, ViTImageProcessor
+from transformers import ViTFeatureExtractor, ViTForImageClassification
 from sentence_transformers import SentenceTransformer
 import matplotlib.pyplot as plt
 import logging
@@ -100,24 +100,48 @@ class ImageAnalyzer:
         self.defect_classes = ["spalling", "structural_cracks", "surface_deterioration"]
         
         try:
+            # Use feature extractor instead of processor
+            self.feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
             self.model = ViTForImageClassification.from_pretrained(
                 "google/vit-base-patch16-224",
-                num_labels=len(self.defect_classes)
+                num_labels=len(self.defect_classes),
+                ignore_mismatched_sizes=True
             ).to(self.device)
-            self.processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
+            
+            # 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,
+                    out_features=len(self.defect_classes)
+                )
+                
         except Exception as e:
             logger.error(f"Model initialization error: {e}")
             self.model = None
-            self.processor = None
+            self.feature_extractor = None
+
+    def preprocess_image(self, image):
+        """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
 
     def analyze_image(self, image):
+        """Analyze image for defects"""
         try:
-            # Ensure image is RGB
-            if image.mode != 'RGB':
-                image = image.convert('RGB')
+            # Preprocess image
+            processed_image = self.preprocess_image(image)
             
-            # Process image
-            inputs = self.processor(images=image, return_tensors="pt")
+            # Extract features
+            inputs = self.feature_extractor(
+                images=processed_image,
+                return_tensors="pt"
+            )
             inputs = {k: v.to(self.device) for k, v in inputs.items()}
             
             # Get predictions
@@ -126,10 +150,24 @@ class ImageAnalyzer:
             
             # Get probabilities
             probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
-            return {self.defect_classes[i]: float(probs[i]) for i in range(len(self.defect_classes))}
+            
+            # Add confidence threshold
+            confidence_threshold = 0.3
+            results = {
+                self.defect_classes[i]: float(probs[i]) 
+                for i in range(len(self.defect_classes))
+                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])}
+            
+            return results
             
         except Exception as e:
-            logger.error(f"Analysis error: {e}")
+            logger.error(f"Analysis error: {str(e)}")
             return None
 
 def get_groq_response(query: str, context: str) -> str:
@@ -170,11 +208,15 @@ def main():
     
     st.title("🏗️ Construction Defect 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()
+    # 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
     
     # Create two columns
     col1, col2 = st.columns([1, 1])
@@ -241,7 +283,7 @@ def main():
                 with st.expander("View retrieved information"):
                     st.text(context)
 
-    # Sidebar for information
+    # Sidebar for information and settings
     with st.sidebar:
         st.header("About")
         st.write("""
@@ -259,6 +301,12 @@ def main():
             st.success("Groq API: Connected")
         else:
             st.error("Groq API: Not configured")
+        
+        # Add settings section
+        st.subheader("Settings")
+        if st.button("Clear Session"):
+            st.session_state.clear()
+            st.success("Session cleared!")
 
 if __name__ == "__main__":
     main()