Update app.py

app.py

@@ -4,9 +4,13 @@ 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')
+
 # Constants
 MAX_FILE_SIZE = 5 * 1024 * 1024  # 5MB
 MAX_IMAGE_SIZE = 1024  # Maximum dimension for images
@@ -16,317 +20,146 @@ MODEL = None
 PROCESSOR = None
 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()
 
-def init_session_state():
-    """Initialize session state variables"""
-    if 'history' not in st.session_state:
-        st.session_state.history = []
-    if 'dark_mode' not in st.session_state:
-        st.session_state.dark_mode = False
-
+# Session state initialization
 @st.cache_resource(show_spinner="Loading AI model...")
 def load_model():
     """Load and cache the model and processor"""
     try:
         model_name = "google/vit-base-patch16-224"
-        # Initialize the processor first
         processor = ViTImageProcessor.from_pretrained(model_name)
-        
-        # Load model with specific device configuration
         device = "cuda" if torch.cuda.is_available() else "cpu"
+
         model = ViTForImageClassification.from_pretrained(
             model_name,
             num_labels=len(DAMAGE_TYPES),
             ignore_mismatched_sizes=True,
         ).to(device)
-        
-        # Ensure model is in evaluation mode
+
         model.eval()
-        
+        logging.info("Model loaded successfully.")
         return model, processor
     except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
+        logging.error(f"Failed to load model: {str(e)}")
+        st.error("Error loading model. Please restart the app.")
         return None, None
 
+# Initialize RAG system
+@st.cache_resource
 def init_rag_system():
-    """Initialize the RAG system with knowledge base"""
     global RAG_SYSTEM
-    if RAG_SYSTEM is None:
+    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.")
 
+# Image validation
 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.")
+        st.warning("Large image detected. Resizing for better performance.")
     if image.format not in ['JPEG', 'PNG']:
-        st.warning("Image format not optimal. Consider using JPEG or PNG for better performance.")
+        st.warning("Non-optimal image format. Use JPEG or PNG.")
 
+# Image preprocessing
 def preprocess_image(uploaded_file):
-    """Preprocess and validate uploaded image"""
     try:
         image = Image.open(uploaded_file)
-        # Resize if image is too large
         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)}")
+        logging.error(f"Error processing image: {str(e)}")
+        st.error("Image processing error.")
         return None
 
+# Damage analysis
 def analyze_damage(image, model, processor):
-    """Analyze structural damage in the image"""
     try:
         device = next(model.parameters()).device
         with torch.no_grad():
             image = image.convert('RGB')
             inputs = processor(images=image, return_tensors="pt")
-            # Move inputs to the same device as model
             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()  # Move results back to CPU
-    except RuntimeError as e:
-        if "out of memory" in str(e):
-            cleanup_memory()
-            st.error("Out of memory. Please try with a smaller image.")
-        else:
-            st.error(f"Error analyzing image: {str(e)}")
+            return probs.cpu()
+    except Exception as e:
+        logging.error(f"Error analyzing image: {str(e)}")
+        st.error("Image analysis failed.")
         return None
 
-def get_custom_css():
-    """Return custom CSS styles"""
-    return """
-    <style>
-    .main {
-        padding: 2rem;
-    }
-    .stProgress > div > div > div > div {
-        background-image: linear-gradient(to right, var(--progress-color, #ff6b6b), var(--progress-color-end, #f06595));
-    }
-    .damage-card {
-        padding: 1.5rem;
-        border-radius: 0.5rem;
-        background: var(--card-bg, #f8f9fa);
-        margin-bottom: 1rem;
-        border: 1px solid var(--border-color, #dee2e6);
-        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
-    }
-    .damage-header {
-        font-size: 1.25rem;
-        font-weight: bold;
-        margin-bottom: 1rem;
-        color: var(--text-color, #212529);
-    }
-    .dark-mode {
-        background-color: #1a1a1a;
-        color: #ffffff;
-    }
-    .dark-mode .damage-card {
-        background: #2d2d2d;
-        border-color: #404040;
-    }
-    </style>
-    """
-
-def display_header():
-    """Display application header"""
-    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
-    )
-
+# Display enhanced analysis
 def display_enhanced_analysis(damage_type, confidence):
-    """Display enhanced analysis from RAG system"""
     try:
         enhanced_info = RAG_SYSTEM.get_enhanced_analysis(damage_type, confidence)
-        
         st.markdown("### 🔍 Enhanced Analysis")
         
         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)
-                
-        custom_query = st.text_input(
-            "Ask specific questions about this damage type:",
-            placeholder="E.g., What are the long-term implications of this damage?",
-            key=f"text_input_{damage_type}"
-        )
-        if custom_query:
-            custom_results = RAG_SYSTEM.get_enhanced_analysis(
-                damage_type,
-                confidence,
-                custom_query=custom_query
-            )
-            st.markdown("### 💡 Custom Query Results")
-            for category, results in custom_results.items():
-                if results:
-                    st.markdown(f"**{category.replace('_', ' ').title()}:**")
-                    for result in results:
-                        st.markdown(result)
     except Exception as e:
-        st.error(f"Error generating enhanced analysis: {str(e)}")
-
-def display_analysis_results(predictions, analysis_time):
-    """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']
-            cases = KNOWLEDGE_BASE[damage_type]
-            
-            with st.expander(f"{damage_type.replace('_', ' ').title()} - {confidence:.1f}%", expanded=True):
-                st.markdown(
-                    f"""
-                    <style>
-                    .stProgress > div > div > div > div {{
-                        background-color: {DAMAGE_TYPES[idx]['color']} !important;
-                    }}
-                    </style>
-                    """,
-                    unsafe_allow_html=True
-                )
-                st.progress(confidence / 100)
-                
-                tabs = st.tabs(["📋 Details", "🔧 Repairs", "⚠️ Actions"])
-                
-                with tabs[0]:
-                    for case in cases:
-                        st.markdown(f"""
-                        - **Severity:** {case['severity']}
-                        - **Description:** {case['description']}
-                        - **Location:** {case['location']}
-                        - **Required Expertise:** {case['required_expertise']}
-                        """)
-                
-                with tabs[1]:
-                    for step in cases[0]['repair_method']:
-                        st.markdown(f"✓ {step}")
-                    st.info(f"**Estimated Cost:** {cases[0]['estimated_cost']}")
-                    st.info(f"**Timeframe:** {cases[0]['timeframe']}")
-                
-                with tabs[2]:
-                    st.warning("**Immediate Actions Required:**")
-                    st.markdown(cases[0]['immediate_action'])
-                    st.success("**Prevention Measures:**")
-                    st.markdown(cases[0]['prevention'])
-            
-            # Display enhanced analysis
-            display_enhanced_analysis(damage_type, confidence)
-    
-    if not detected:
-        st.info("No significant structural damage detected. Regular maintenance recommended.")
+        logging.error(f"Failed to generate enhanced analysis: {str(e)}")
+        st.error("Error generating enhanced analysis.")
 
+# Main function
 def main():
-    """Main application function"""
-    init_session_state()
     st.set_page_config(
         page_title="Structural Damage Analyzer Pro",
         page_icon="🏗️",
-        layout="wide",
-        initial_sidebar_state="expanded"
+        layout="wide"
     )
-
-    st.markdown(get_custom_css(), 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}")
-    
-    display_header()
+    st.title("🏗️ Structural Damage Analyzer Pro")
 
     # Load model and initialize RAG system
     global MODEL, PROCESSOR
     if MODEL is None or PROCESSOR is None:
-        with st.spinner("Loading AI model..."):
-            MODEL, PROCESSOR = load_model()
-            if MODEL is None:
-                st.error("Failed to load model. Please refresh the page.")
-                return
-    
+        MODEL, PROCESSOR = load_model()
     init_rag_system()
 
-    # File upload
     uploaded_file = st.file_uploader(
-        "Drag and drop or click to upload an image",
-        type=['jpg', 'jpeg', 'png'],
-        help="Supported formats: JPG, JPEG, PNG"
+        "Upload an image for analysis (JPG, PNG)",
+        type=['jpg', 'jpeg', 'png']
     )
 
     if 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)
-            
-            col1, col2 = st.columns([1, 1])
-            
-            with col1:
-                st.image(image, caption="Uploaded Structure", use_container_width=True)
-            
-            with col2:
-                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:
-                        display_analysis_results(predictions, analysis_time)
-                        st.session_state.history.append(f"Analyzed image: {uploaded_file.name}")
+        if uploaded_file.size > MAX_FILE_SIZE:
+            st.error("File too large. Limit: 5MB.")
+            return
 
-        except Exception as e:
-            cleanup_memory()
-            st.error(f"Error processing image: {str(e)}")
-            st.info("Please try uploading a different image.")
+        image = preprocess_image(uploaded_file)
+        validate_image(image)
 
-    # Footer
-    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
-    )
+        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)
 
 if __name__ == "__main__":
-    main()
+    main()