Update app.py

app.py

@@ -1,13 +1,39 @@
+# app.py
 import streamlit as st
 from transformers import ViTForImageClassification, ViTImageProcessor
 from PIL import Image
 import torch
 import time
 import gc
+import logging
+from knowledge_base import KNOWLEDGE_BASE, DAMAGE_TYPES
+from rag_utils import RAGSystem
+import structlog
+from typing import Optional, Dict, Any
+from functools import lru_cache
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = structlog.get_logger()
 
 # Constants
 MAX_FILE_SIZE = 5 * 1024 * 1024  # 5MB
 MAX_IMAGE_SIZE = 1024  # Maximum dimension for images
+MODEL = None
+PROCESSOR = None
+RAG_SYSTEM = None
+
+def handle_exceptions(func):
+    """Decorator for exception handling"""
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            cleanup_memory()
+            st.error(f"Error in {func.__name__}: {str(e)}")
+            logger.error(f"Error in {func.__name__}: {str(e)}", exc_info=True)
+            return None
+    return wrapper
 
 def cleanup_memory():
     """Clean up memory and GPU cache"""
@@ -22,7 +48,7 @@ def init_session_state():
     if 'dark_mode' not in st.session_state:
         st.session_state.dark_mode = False
 
-@st.cache_resource(show_spinner="Loading AI model...")
+@st.cache_resource(show_spinner="Loading AI model...", ttl=3600*24)
 def load_model():
     """Load and cache the model and processor"""
     try:
@@ -35,43 +61,71 @@ def load_model():
             ignore_mismatched_sizes=True,
         ).to(device)
         model.eval()
+        logger.info("Model loaded successfully", device=device)
         return model, processor
     except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
+        logger.error("Error loading model", error=str(e))
         return None, None
 
-def validate_image(image):
-    """Validate image size and format"""
-    if image.size[0] * image.size[1] > 1024 * 1024:
-        st.warning("Large image detected. The image will be resized for better performance.")
-    if image.format not in ['JPEG', 'PNG']:
-        st.warning("Image format not optimal. Consider using JPEG or PNG for better performance.")
-
-def preprocess_image(uploaded_file):
+def validate_upload(file) -> bool:
+    """Validate uploaded file for security"""
+    if not file:
+        return False
+    
+    allowed_extensions = {'jpg', 'jpeg', 'png'}
+    if not file.name.lower().endswith(tuple(allowed_extensions)):
+        st.error("Invalid file type. Please upload a JPG or PNG image.")
+        return False
+    
+    if file.size > MAX_FILE_SIZE:
+        st.error("File too large. Maximum size is 5MB.")
+        return False
+    
+    if file.type not in ['image/jpeg', 'image/png']:
+        st.error("Invalid file content type.")
+        return False
+    
+    return True
+
+@handle_exceptions
+def preprocess_image(uploaded_file) -> Optional[Image.Image]:
     """Preprocess and validate uploaded image"""
-    try:
-        image = Image.open(uploaded_file)
-        if max(image.size) > MAX_IMAGE_SIZE:
-            ratio = MAX_IMAGE_SIZE / max(image.size)
-            new_size = tuple([int(dim * ratio) for dim in image.size])
-            image = image.resize(new_size, Image.Resampling.LANCZOS)
-        return image
-    except Exception as e:
-        st.error(f"Error processing image: {str(e)}")
-        return None
-
-def analyze_damage(image, model, processor):
+    image = Image.open(uploaded_file)
+    if max(image.size) > MAX_IMAGE_SIZE:
+        ratio = MAX_IMAGE_SIZE / max(image.size)
+        new_size = tuple([int(dim * ratio) for dim in image.size])
+        image = image.resize(new_size, Image.Resampling.LANCZOS)
+    return image
+
+@handle_exceptions
+def analyze_damage(image: Image.Image, model: ViTForImageClassification, 
+                  processor: ViTImageProcessor) -> Optional[torch.Tensor]:
     """Analyze structural damage in the image"""
+    progress_bar = st.progress(0)
+    stages = ['Preprocessing', 'Analysis', 'Results Generation']
+    
     try:
         device = next(model.parameters()).device
-        with torch.no_grad():
-            image = image.convert('RGB')
-            inputs = processor(images=image, return_tensors="pt")
-            inputs = {k: v.to(device) for k, v in inputs.items()}
-            outputs = model(**inputs)
-            probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
-            cleanup_memory()
-            return probs.cpu()
+        
+        for i, stage in enumerate(stages):
+            progress_bar.progress((i + 1) / len(stages))
+            st.write(f"Stage {i+1}/{len(stages)}: {stage}")
+            
+            if i == 0:  # Preprocessing
+                image = image.convert('RGB')
+                inputs = processor(images=image, return_tensors="pt")
+                inputs = {k: v.to(device) for k, v in inputs.items()}
+            
+            elif i == 1:  # Analysis
+                with torch.no_grad():
+                    outputs = model(**inputs)
+                    probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
+            
+            elif i == 2:  # Results Generation
+                result = probs.cpu()
+                
+        return result
+        
     except RuntimeError as e:
         if "out of memory" in str(e):
             cleanup_memory()
@@ -80,78 +134,120 @@ def analyze_damage(image, model, processor):
             st.error(f"Error analyzing image: {str(e)}")
         return None
 
+def display_analysis_results(predictions: torch.Tensor, analysis_time: float):
+    """Display analysis results with damage details"""
+    st.markdown("### 📊 Analysis Results")
+    st.markdown(f"*Analysis completed in {analysis_time:.2f} seconds*")
+    
+    detected = False
+    for idx, prob in enumerate(predictions):
+        confidence = float(prob) * 100
+        if confidence > 15:
+            detected = True
+            damage_type = DAMAGE_TYPES[idx]['name']
+            
+            with st.expander(f"{damage_type.replace('_', ' ').title()} - {confidence:.1f}%", expanded=True):
+                st.progress(confidence / 100)
+                
+                # Get enhanced analysis from RAG system
+                analysis = RAG_SYSTEM.get_enhanced_analysis(damage_type, confidence)
+                
+                tabs = st.tabs(["📋 Details", "🔧 Repairs", "⚠️ Safety"])
+                
+                with tabs[0]:
+                    for detail in analysis['technical_details']:
+                        st.markdown(detail)
+                
+                with tabs[1]:
+                    for rec in analysis['expert_recommendations']:
+                        st.markdown(rec)
+                
+                with tabs[2]:
+                    for safety in analysis['safety_considerations']:
+                        st.warning(safety)
+    
+    if not detected:
+        st.info("No significant structural damage detected. Regular maintenance recommended.")
+
 def main():
+    """Main application function"""
     st.set_page_config(
         page_title="Structural Damage Analyzer Pro",
         page_icon="🏗️",
         layout="wide",
         initial_sidebar_state="expanded"
     )
-
-    # Initialize session state
+    
     init_session_state()
-
-    # Display header
-    st.markdown(
-        """
+    
+    st.markdown("""
         <div style='text-align: center; padding: 1rem;'>
             <h1>🏗️ Structural Damage Analyzer Pro</h1>
             <p style='font-size: 1.2rem;'>Advanced AI-powered structural damage assessment tool</p>
         </div>
-        """,
-        unsafe_allow_html=True
-    )
-
-    # Load model
-    model, processor = load_model()
-    if model is None:
-        st.error("Failed to load model. Please refresh the page.")
-        return
-
+    """, unsafe_allow_html=True)
+    
+    # Sidebar
+    with st.sidebar:
+        st.markdown("### ⚙️ Settings")
+        st.session_state.dark_mode = st.toggle("Dark Mode", st.session_state.dark_mode)
+        st.markdown("### 📖 Analysis History")
+        if st.session_state.history:
+            for item in st.session_state.history[-5:]:
+                st.markdown(f"- {item}")
+    
+    # Load model and initialize RAG system
+    global MODEL, PROCESSOR, RAG_SYSTEM
+    if MODEL is None or PROCESSOR is None:
+        MODEL, PROCESSOR = load_model()
+        if MODEL is None:
+            st.error("Failed to load model. Please refresh the page.")
+            return
+    
+    if RAG_SYSTEM is None:
+        RAG_SYSTEM = RAGSystem()
+        RAG_SYSTEM.initialize_knowledge_base(KNOWLEDGE_BASE)
+    
     # File upload
     uploaded_file = st.file_uploader(
         "Upload an image for analysis",
         type=['jpg', 'jpeg', 'png'],
         help="Supported formats: JPG, JPEG, PNG"
     )
-
-    if uploaded_file:
+    
+    if uploaded_file and validate_upload(uploaded_file):
         try:
-            if uploaded_file.size > MAX_FILE_SIZE:
-                st.error("File size too large. Please upload an image smaller than 5MB.")
-                return
-
             image = preprocess_image(uploaded_file)
             if image is None:
                 return
-
-            validate_image(image)
             
-            # Display image and analyze
-            st.image(image, caption="Uploaded Structure", use_column_width=True)
+            col1, col2 = st.columns([1, 1])
+            
+            with col1:
+                st.image(image, caption="Uploaded Structure", use_column_width=True)
             
-            with st.spinner("🔍 Analyzing damage..."):
-                predictions = analyze_damage(image, model, processor)
+            with col2:
+                start_time = time.time()
+                predictions = analyze_damage(image, MODEL, PROCESSOR)
+                
                 if predictions is not None:
-                    st.success("Analysis complete!")
-                    # Add analysis display logic here based on your DAMAGE_TYPES
-
+                    analysis_time = time.time() - start_time
+                    display_analysis_results(predictions, analysis_time)
+                    st.session_state.history.append(f"Analyzed {uploaded_file.name}")
+        
         except Exception as e:
+            logger.error("Error in main processing loop", error=str(e))
             cleanup_memory()
-            st.error(f"Error processing image: {str(e)}")
-            st.info("Please try uploading a different image.")
-
+            st.error("An error occurred during processing. Please try again.")
+    
     # Footer
     st.markdown("---")
-    st.markdown(
-        """
+    st.markdown("""
         <div style='text-align: center'>
             <p>🏗️ Structural Damage Analyzer Pro | Built with Streamlit & Transformers</p>
             <p style='font-size: 0.8rem;'>For professional use only. Always consult with a structural engineer.</p>
         </div>
-        """,
-        unsafe_allow_html=True
-    )
+    """, unsafe_allow_html=True)
 
 if __name__ == "__main__":
     main()