app.py

import streamlit as st
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"]
        
        try:
            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"Model initialization error: {e}")
            self.model = None
            self.processor = None

    def analyze_image(self, image):
        try:
            # Ensure image is RGB
            if image.mode != 'RGB':
                image = image.convert('RGB')
            
            # Process image
            inputs = self.processor(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]
            return {self.defect_classes[i]: float(probs[i]) for i in range(len(self.defect_classes))}
            
        except Exception as 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 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.set_page_config(
        page_title="Construction Defect Analyzer",
        page_icon="🏗️",
        layout="wide"
    )
    
    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()
    
    # 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:
                # 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:
                        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.")
                    
            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()