Update app.py

app.py

@@ -4,162 +4,270 @@ 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
-
-# Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+import os
 
 # Constants
 MAX_FILE_SIZE = 5 * 1024 * 1024  # 5MB
 MAX_IMAGE_SIZE = 1024  # Maximum dimension for images
+MODEL_NAME = "google/vit-base-patch16-224"
+CACHE_DIR = "/tmp/model_cache"  # HF Spaces compatible cache directory
+
+# Ensure cache directory exists
+os.makedirs(CACHE_DIR, exist_ok=True)
 
-# Cache the model and RAG system globally
-MODEL = None
-PROCESSOR = None
-RAG_SYSTEM = None
+# Initialize session state for caching
+if 'model' not in st.session_state:
+    st.session_state.model = None
+if 'processor' not in st.session_state:
+    st.session_state.processor = None
+if 'rag_system' not in st.session_state:
+    st.session_state.rag_system = None
 
-# Cleanup function for memory
 def cleanup_memory():
     """Clean up memory and GPU cache"""
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
 
-# Session state initialization
 @st.cache_resource(show_spinner="Loading AI model...")
 def load_model():
-    """Load and cache the model and processor"""
+    """Load and cache the model and processor with error handling"""
     try:
-        model_name = "google/vit-base-patch16-224"
-        processor = ViTImageProcessor.from_pretrained(model_name)
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-
+        # Initialize processor with cache directory
+        processor = ViTImageProcessor.from_pretrained(
+            MODEL_NAME,
+            cache_dir=CACHE_DIR,
+            local_files_only=False
+        )
+        
+        # Determine device - prefer CPU on Hugging Face Spaces
+        device = "cpu"  # Default to CPU for stability
+        
+        # Load model with specific configuration
         model = ViTForImageClassification.from_pretrained(
-            model_name,
+            MODEL_NAME,
             num_labels=len(DAMAGE_TYPES),
             ignore_mismatched_sizes=True,
+            cache_dir=CACHE_DIR,
+            local_files_only=False
         ).to(device)
-
-        model.eval()
-        logging.info("Model loaded successfully.")
+        
+        model.eval()  # Set to evaluation mode
         return model, processor
     except Exception as e:
-        logging.error(f"Failed to load model: {str(e)}")
-        st.error("Error loading model. Please restart the app.")
+        st.error(f"Error loading model: {str(e)}")
+        st.info("Attempting to reload model... Please wait.")
+        cleanup_memory()
         return None, None
 
-# Initialize RAG system
-@st.cache_resource
 def init_rag_system():
-    global RAG_SYSTEM
-    try:
-        RAG_SYSTEM = RAGSystem()
-        RAG_SYSTEM.initialize_knowledge_base(KNOWLEDGE_BASE)
-        logging.info("RAG system initialized successfully.")
-    except Exception as e:
-        logging.error(f"Failed to initialize RAG system: {str(e)}")
-        st.error("Error initializing knowledge base.")
+    """Initialize RAG system with error handling"""
+    if st.session_state.rag_system is None:
+        try:
+            st.session_state.rag_system = RAGSystem()
+            st.session_state.rag_system.initialize_knowledge_base(KNOWLEDGE_BASE)
+        except Exception as e:
+            st.error(f"Error initializing RAG system: {str(e)}")
+            st.session_state.rag_system = None
 
-# Image validation
-def validate_image(image):
-    if image.size[0] * image.size[1] > 1024 * 1024:
-        st.warning("Large image detected. Resizing for better performance.")
-    if image.format not in ['JPEG', 'PNG']:
-        st.warning("Non-optimal image format. Use JPEG or PNG.")
-
-# Image preprocessing
-def preprocess_image(uploaded_file):
+def process_image(image):
+    """Process and validate image with enhanced error handling"""
     try:
-        image = Image.open(uploaded_file)
+        # Convert to RGB if necessary
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        # Resize if needed
         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:
-        logging.error(f"Error processing image: {str(e)}")
-        st.error("Image processing error.")
+        st.error(f"Error processing image: {str(e)}")
         return None
 
-# Damage analysis
 def analyze_damage(image, model, processor):
+    """Analyze structural damage with enhanced error handling and memory management"""
     try:
         device = next(model.parameters()).device
         with torch.no_grad():
-            image = image.convert('RGB')
+            # Process image
             inputs = processor(images=image, return_tensors="pt")
             inputs = {k: v.to(device) for k, v in inputs.items()}
+            
+            # Run inference
             outputs = model(**inputs)
             probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
+            
+            # Clean up
             cleanup_memory()
             return probs.cpu()
+    except RuntimeError as e:
+        if "out of memory" in str(e):
+            cleanup_memory()
+            st.error("Memory error. Processing with reduced image size...")
+            # Retry with smaller image
+            image = image.resize((224, 224), Image.Resampling.LANCZOS)
+            return analyze_damage(image, model, processor)
+        else:
+            st.error(f"Error during analysis: {str(e)}")
+        return None
     except Exception as e:
-        logging.error(f"Error analyzing image: {str(e)}")
-        st.error("Image analysis failed.")
+        st.error(f"Unexpected error: {str(e)}")
         return None
 
-# Display enhanced analysis
-def display_enhanced_analysis(damage_type, confidence):
+def display_analysis_results(predictions, analysis_time):
+    """Display analysis results with enhanced visualization and error handling"""
     try:
-        enhanced_info = RAG_SYSTEM.get_enhanced_analysis(damage_type, confidence)
-        st.markdown("### 🔍 Enhanced Analysis")
+        st.markdown("### 📊 Analysis Results")
+        st.markdown(f"*Analysis completed in {analysis_time:.2f} seconds*")
         
-        with st.expander("📚 Technical Details", expanded=True):
-            for detail in enhanced_info["technical_details"]:
-                st.markdown(detail)
-
-        with st.expander("⚠️ Safety Considerations"):
-            for safety in enhanced_info["safety_considerations"]:
-                st.warning(safety)
-
-        with st.expander("👷 Expert Recommendations"):
-            for rec in enhanced_info["expert_recommendations"]:
-                st.info(rec)
+        detected = False
+        for idx, prob in enumerate(predictions):
+            confidence = float(prob) * 100
+            if confidence > 15:  # Threshold for displaying results
+                detected = True
+                damage_type = DAMAGE_TYPES[idx]['name']
+                risk_level = DAMAGE_TYPES[idx]['risk']
+                
+                # Create expander with color-coded header
+                with st.expander(
+                    f"🔍 {damage_type.replace('_', ' ').title()} - {confidence:.1f}% ({risk_level})",
+                    expanded=True
+                ):
+                    # Display confidence bar
+                    st.progress(confidence / 100)
+                    
+                    # Create tabs for organized information
+                    details_tab, repair_tab, action_tab = st.tabs([
+                        "📋 Details", "🔧 Repair Plan", "⚠️ Actions Needed"
+                    ])
+                    
+                    with details_tab:
+                        display_damage_details(damage_type, confidence)
+                    
+                    with repair_tab:
+                        display_repair_plan(damage_type)
+                    
+                    with action_tab:
+                        display_action_items(damage_type)
+                    
+                    # Display enhanced analysis if RAG system is available
+                    if st.session_state.rag_system:
+                        display_enhanced_analysis(damage_type, confidence)
+        
+        if not detected:
+            st.success("No significant structural damage detected. Regular maintenance recommended.")
+            
     except Exception as e:
-        logging.error(f"Failed to generate enhanced analysis: {str(e)}")
-        st.error("Error generating enhanced analysis.")
+        st.error(f"Error displaying results: {str(e)}")
 
-# Main function
 def main():
-    st.set_page_config(
-        page_title="Structural Damage Analyzer Pro",
-        page_icon="🏗️",
-        layout="wide"
-    )
-    st.title("🏗️ Structural Damage Analyzer Pro")
-
-    # Load model and initialize RAG system
-    global MODEL, PROCESSOR
-    if MODEL is None or PROCESSOR is None:
-        MODEL, PROCESSOR = load_model()
-    init_rag_system()
-
-    uploaded_file = st.file_uploader(
-        "Upload an image for analysis (JPG, PNG)",
-        type=['jpg', 'jpeg', 'png']
-    )
+    """Main application function with enhanced error handling and UI"""
+    try:
+        # Page configuration
+        st.set_page_config(
+            page_title="Structural Damage Analyzer Pro",
+            page_icon="🏗️",
+            layout="wide",
+            initial_sidebar_state="expanded"
+        )
+        
+        # Custom CSS
+        st.markdown(get_custom_css(), unsafe_allow_html=True)
+        
+        # Header
+        display_header()
+        
+        # Initialize systems
+        if st.session_state.model is None or st.session_state.processor is None:
+            with st.spinner("Initializing AI model..."):
+                model, processor = load_model()
+                if model is None:
+                    st.error("Failed to initialize model. Please refresh the page.")
+                    return
+                st.session_state.model = model
+                st.session_state.processor = processor
+        
+        init_rag_system()
+        
+        # File upload section
+        uploaded_file = st.file_uploader(
+            "Upload structural image for analysis",
+            type=['jpg', 'jpeg', 'png'],
+            help="Maximum file size: 5MB"
+        )
+        
+        if uploaded_file:
+            process_uploaded_file(uploaded_file)
+        
+        # Footer
+        display_footer()
+        
+    except Exception as e:
+        st.error(f"Application error: {str(e)}")
+        st.info("Please refresh the page and try again.")
+        cleanup_memory()
 
-    if uploaded_file:
+def process_uploaded_file(uploaded_file):
+    """Process uploaded file with comprehensive error handling"""
+    try:
+        # Validate file size
         if uploaded_file.size > MAX_FILE_SIZE:
-            st.error("File too large. Limit: 5MB.")
+            st.error("File too large. Please upload an image smaller than 5MB.")
             return
+        
+        # Process image
+        image = Image.open(uploaded_file)
+        processed_image = process_image(image)
+        if processed_image is None:
+            return
+        
+        # Display layout
+        col1, col2 = st.columns([1, 1])
+        with col1:
+            st.image(processed_image, caption="Uploaded Structure", use_column_width=True)
+        
+        with col2:
+            with st.spinner("🔍 Analyzing structural damage..."):
+                start_time = time.time()
+                predictions = analyze_damage(
+                    processed_image,
+                    st.session_state.model,
+                    st.session_state.processor
+                )
+                if predictions is not None:
+                    analysis_time = time.time() - start_time
+                    display_analysis_results(predictions, analysis_time)
+                
+    except Exception as e:
+        st.error(f"Error processing upload: {str(e)}")
+        cleanup_memory()
 
-        image = preprocess_image(uploaded_file)
-        validate_image(image)
-
-        st.image(image, caption="Uploaded Image", use_column_width=True)
-
-        with st.spinner("Analyzing damage..."):
-            start_time = time.time()
-            predictions = analyze_damage(image, MODEL, PROCESSOR)
-            analysis_time = time.time() - start_time
-
-            if predictions is not None:
-                st.markdown(f"*Analysis completed in {analysis_time:.2f} seconds*")
-                confidence = float(predictions[0]) * 100
-                display_enhanced_analysis(DAMAGE_TYPES[0]['name'], confidence)
+def get_custom_css():
+    """Return custom CSS for enhanced UI"""
+    return """
+    <style>
+    .main {
+        padding: 1rem;
+    }
+    .stProgress > div > div > div > div {
+        background-image: linear-gradient(to right, #ff6b6b, #f06595);
+    }
+    .damage-card {
+        padding: 1rem;
+        border-radius: 0.5rem;
+        background: var(--background-color, #ffffff);
+        margin-bottom: 1rem;
+        border: 1px solid var(--border-color, #e0e0e0);
+        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+    }
+    </style>
+    """
 
 if __name__ == "__main__":
-    main()
+    main()