Update app.py

app.py

@@ -1,313 +1,98 @@
 import streamlit as st
-import os
-from groq import Groq
-from transformers import ViTForImageClassification, ViTImageProcessor
-from sentence_transformers import SentenceTransformer
-from PIL import Image
 import torch
+from PIL import Image
 import numpy as np
-from typing import List, Dict, Tuple
-import faiss
-import json
-import cv2
-import logging
-from datetime import datetime
+from transformers import ViTForImageClassification, ViTImageProcessor
 import matplotlib.pyplot as plt
+import logging
 
 # Setup logging
 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:
+class ImageAnalyzer:
     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"
-                }
-            ],
-            "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}})
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.defect_classes = ["spalling", "structural_cracks", "surface_deterioration"]
         
-        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"""
+        # Load model and processor
         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
+            self.model = ViTForImageClassification.from_pretrained(
+                "google/vit-base-patch16-224",
+                num_labels=len(self.defect_classes)
+            ).to(self.device)
+            self.processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
         except Exception as e:
-            logger.error(f"Error retrieving context: {e}")
-            return ""
-
-class ImageAnalyzer:
-    def __init__(self):
-        self.device = DEVICE
-        self.defect_classes = ["spalling", "structural_cracks", "surface_deterioration"]
-        self.model = load_vit_model()
-        self.processor = load_vit_processor()
-        self.history = []
+            logger.error(f"Model initialization error: {e}")
+            self.model = None
+            self.processor = None
 
-    def preprocess_image(self, image: Image.Image) -> torch.Tensor:
-        """Preprocess image for model input"""
+    def analyze_image(self, image):
         try:
             # Ensure image is RGB
             if image.mode != 'RGB':
                 image = image.convert('RGB')
             
-            # Process image using ViT processor
+            # Process image
             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.preprocess_image(image)
-            if inputs is None:
-                raise ValueError("Image preprocessing failed")
-
-            # Get model predictions
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+            
+            # Get predictions
             with torch.no_grad():
                 outputs = self.model(**inputs)
             
-            # 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))
-            }
-            
-            # Basic image statistics
-            img_array = np.array(image)
-            stats = {
-                "mean_brightness": float(np.mean(img_array)),
-                "image_size": f"{image.size[0]}x{image.size[1]}"
-            }
-            
-            result = {
-                "defect_probabilities": defect_probs,
-                "image_statistics": stats,
-                "timestamp": datetime.now().isoformat()
-            }
-            
-            self.history.append(result)
-            return result
+            # 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))}
             
         except Exception as e:
-            logger.error(f"Image analysis error: {e}")
+            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 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="Construction Defect Analyzer",
-        page_icon="🏗️",
-        layout="wide"
-    )
-
-    st.title("🏗️ Construction Defect Analyzer")
-
-    # Initialize systems in session state
-    if 'rag_system' not in st.session_state:
-        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'],
-            key="image_uploader"
-        )
-        
-        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)
+    st.title("Construction Defect Analyzer")
+    
+    # Initialize analyzer
+    if 'analyzer' not in st.session_state:
+        st.session_state.analyzer = ImageAnalyzer()
+    
+    # File uploader
+    st.write("Upload a construction image for analysis")
+    uploaded_file = st.file_uploader("Choose an image...", 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 results:
-                    status_text.text("Analysis complete!")
-                    progress_bar.progress(100)
+                    st.success('Analysis complete!')
                     
-                    st.subheader("Detected Defects")
+                    # Display results
+                    st.subheader("Defect Probabilities")
                     
                     # Create bar chart
-                    defect_probs = results["defect_probabilities"]
-                    fig, ax = plt.subplots(figsize=(8, 4))
-                    defects = list(defect_probs.keys())
-                    probs = list(defect_probs.values())
+                    fig, ax = plt.subplots()
+                    defects = list(results.keys())
+                    probs = list(results.values())
                     ax.barh(defects, probs)
                     ax.set_xlim(0, 1)
-                    ax.set_xlabel("Probability")
+                    plt.tight_layout()
                     st.pyplot(fig)
-                    
-                    # Show image statistics
-                    with st.expander("Image Details"):
-                        st.json(results["image_statistics"])
                 else:
-                    status_text.text("Analysis failed. Please try again.")
-                    progress_bar.empty()
+                    st.error("Analysis failed. Please try again.")
                     
-            except Exception as e:
-                st.error(f"Error processing image: {str(e)}")
-                logger.error(f"Image processing error: {e}")
-
-    with col2:
-        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?"
-        )
-        
-        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 Response")
-                st.write(response)
-                
-                with st.expander("View 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")
+        except Exception as e:
+            st.error(f"Error: {str(e)}")
+            logger.error(f"Process error: {e}")
 
 if __name__ == "__main__":
     main()