Update app.py

app.py

@@ -2,216 +2,322 @@ import streamlit as st
 from transformers import ViTForImageClassification, ViTImageProcessor
 from PIL import Image
 import torch
-import json
+import time
+import io
+import base64
 
-# Embedded knowledge base
+# Cache the model globally
+MODEL = None
+PROCESSOR = None
+
+# Embedded knowledge base (using previous KNOWLEDGE_BASE)
+# Knowledge base
 KNOWLEDGE_BASE = {
     "spalling": [
+        {
+            "severity": "Critical",
+            "description": "Severe concrete spalling with exposed reinforcement and section loss",
+            "repair_method": ["Install temporary support", "Remove deteriorated concrete", "Clean and treat reinforcement", "Apply corrosion inhibitor", "Apply bonding agent", "High-strength repair mortar"],
+            "estimated_cost": "Very High ($15,000+)",
+            "timeframe": "3-4 weeks",
+            "location": "Primary structural elements",
+            "required_expertise": "Structural Engineer + Specialist Contractor",
+            "immediate_action": "Evacuate area, install temporary support, prevent access",
+            "prevention": "Regular inspections, waterproofing, chloride protection"
+        },
         {
             "severity": "High",
-            "description": "Severe concrete spalling with exposed reinforcement",
-            "repair_method": ["Remove damaged concrete", "Clean exposed reinforcement", "Apply rust inhibitor", "Apply bonding agent", "Patch with repair mortar"],
-            "estimated_cost": "High ($5,000-$10,000)",
+            "description": "Surface spalling with visible reinforcement",
+            "repair_method": ["Remove damaged concrete", "Treat reinforcement", "Apply repair mortar", "Surface treatment"],
+            "estimated_cost": "High ($8,000-$15,000)",
             "timeframe": "2-3 weeks",
-            "location": "Column/Beam",
+            "location": "Structural elements",
             "required_expertise": "Structural Engineer",
-            "immediate_action": "Area isolation and temporary support if structural",
-            "prevention": "Regular maintenance, waterproofing, proper concrete cover"
-        },
-        {
-            "severity": "Medium",
-            "description": "Surface spalling without exposed reinforcement",
-            "repair_method": ["Remove loose concrete", "Clean surface", "Apply repair mortar", "Surface treatment"],
-            "estimated_cost": "Medium ($2,000-$5,000)",
-            "timeframe": "1-2 weeks",
-            "location": "Non-structural elements",
-            "required_expertise": "Concrete Repair Specialist",
-            "immediate_action": "Mark affected areas and monitor",
-            "prevention": "Surface coating, regular inspections"
+            "immediate_action": "Area isolation, temporary support assessment",
+            "prevention": "Protective coatings, drainage improvement"
         }
     ],
     "reinforcement_corrosion": [
         {
             "severity": "Critical",
-            "description": "Advanced corrosion with significant section loss",
-            "repair_method": ["Install temporary support", "Remove damaged concrete", "Replace/supplement reinforcement", "Apply corrosion inhibitor", "Reconstruct concrete cover"],
-            "estimated_cost": "Very High ($15,000+)",
-            "timeframe": "3-4 weeks",
-            "location": "Primary structural elements",
-            "required_expertise": "Structural Engineer, Specialist Contractor",
-            "immediate_action": "Immediate area evacuation and temporary support",
-            "prevention": "Waterproofing, crack sealing, chloride protection"
+            "description": "Severe corrosion with >30% section loss",
+            "repair_method": ["Structural support", "Remove concrete", "Replace reinforcement", "Corrosion protection", "Concrete repair"],
+            "estimated_cost": "Critical ($20,000+)",
+            "timeframe": "4-6 weeks",
+            "location": "Load-bearing elements",
+            "required_expertise": "Senior Structural Engineer",
+            "immediate_action": "Immediate evacuation, emergency shoring",
+            "prevention": "Waterproofing, cathodic protection"
         }
     ],
     "structural_crack": [
         {
             "severity": "High",
-            "description": "Load-bearing element cracks >3mm",
-            "repair_method": ["Structural assessment", "Crack injection with epoxy", "External reinforcement if needed", "Monitor crack progression"],
-            "estimated_cost": "High ($8,000-$15,000)",
-            "timeframe": "2-3 weeks",
-            "location": "Load-bearing walls/beams",
+            "description": "Cracks >5mm in structural elements",
+            "repair_method": ["Structural analysis", "Epoxy injection", "Carbon fiber reinforcement", "Crack monitoring"],
+            "estimated_cost": "High ($10,000-$20,000)",
+            "timeframe": "2-4 weeks",
+            "location": "Primary structure",
             "required_expertise": "Structural Engineer",
-            "immediate_action": "Install crack gauges, temporary support",
-            "prevention": "proper design, load management, movement joints"
+            "immediate_action": "Install crack monitors, load restriction",
+            "prevention": "Load management, joint maintenance"
         }
     ],
     "dampness": [
         {
             "severity": "Medium",
-            "description": "Persistent moisture penetration",
-            "repair_method": ["Identify water source", "Install drainage system", "Apply waterproofing membrane", "Improve ventilation"],
-            "estimated_cost": "Medium ($3,000-$7,000)",
+            "description": "Active water penetration with efflorescence",
+            "repair_method": ["Water source identification", "Drainage improvement", "Waterproof membrane", "Ventilation"],
+            "estimated_cost": "Medium ($5,000-$10,000)",
             "timeframe": "1-2 weeks",
-            "location": "Walls, floors",
+            "location": "Various",
             "required_expertise": "Waterproofing Specialist",
-            "immediate_action": "Improve ventilation, dehumidification",
-            "prevention": "Proper drainage, regular maintenance of water systems"
+            "immediate_action": "Dehumidification, efflorescence cleaning",
+            "prevention": "Proper drainage, vapor barriers"
         }
     ],
     "no_damage": [
         {
             "severity": "Low",
-            "description": "No visible structural issues",
-            "repair_method": ["Regular inspection", "Preventive maintenance", "Document condition"],
-            "estimated_cost": "Low ($500-$1,000)",
+            "description": "No significant structural issues",
+            "repair_method": ["Regular inspection", "Preventive maintenance"],
+            "estimated_cost": "Low ($500-$2,000)",
             "timeframe": "1-2 days",
-            "location": "General structure",
+            "location": "General",
             "required_expertise": "Building Inspector",
-            "immediate_action": "Continue regular maintenance",
-            "prevention": "Maintain inspection schedule"
+            "immediate_action": "Continue monitoring",
+            "prevention": "Regular maintenance schedule"
         }
     ]
 }
 
 DAMAGE_TYPES = {
-    0: {'name': 'spalling', 'risk': 'High'},
-    1: {'name': 'reinforcement_corrosion', 'risk': 'Critical'},
-    2: {'name': 'structural_crack', 'risk': 'High'},
-    3: {'name': 'dampness', 'risk': 'Medium'},
-    4: {'name': 'no_damage', 'risk': 'Low'}
+    0: {'name': 'spalling', 'risk': 'High', 'color': '#ff4d4d'},
+    1: {'name': 'reinforcement_corrosion', 'risk': 'Critical', 'color': '#800000'},
+    2: {'name': 'structural_crack', 'risk': 'High', 'color': '#ff6b6b'},
+    3: {'name': 'dampness', 'risk': 'Medium', 'color': '#4dabf7'},
+    4: {'name': 'no_damage', 'risk': 'Low', 'color': '#40c057'}
 }
 
+def init_session_state():
+    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
+
 @st.cache_resource
 def load_model():
-    model = ViTForImageClassification.from_pretrained(
-        "google/vit-base-patch16-224",
-        num_labels=len(DAMAGE_TYPES),
-        ignore_mismatched_sizes=True
-    )
-    processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
-    return model, processor
+    try:
+        model = ViTForImageClassification.from_pretrained(
+            "google/vit-base-patch16-224",
+            num_labels=len(DAMAGE_TYPES),
+            ignore_mismatched_sizes=True
+        )
+        processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
+        return model, processor
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        return None, None
 
 def analyze_damage(image, model, processor):
-    image = image.convert('RGB')
-    inputs = processor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
-    probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
-    return probs
+    try:
+        image = image.convert('RGB')
+        inputs = processor(images=image, return_tensors="pt")
+        outputs = model(**inputs)
+        probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
+        return probs
+    except Exception as e:
+        st.error(f"Error analyzing image: {str(e)}")
+        return None
+
+def get_custom_css():
+    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;
+    }
+    .tooltip {
+        position: relative;
+        display: inline-block;
+        border-bottom: 1px dotted #ccc;
+    }
+    .tooltip .tooltiptext {
+        visibility: hidden;
+        background-color: #555;
+        color: #fff;
+        text-align: center;
+        border-radius: 6px;
+        padding: 5px;
+        position: absolute;
+        z-index: 1;
+        bottom: 125%;
+        left: 50%;
+        margin-left: -60px;
+        opacity: 0;
+        transition: opacity 0.3s;
+    }
+    </style>
+    """
+
+def display_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
+    )
 
 def main():
+    init_session_state()
     st.set_page_config(
-        page_title="Structural Damage Analyzer",
+        page_title="Structural Damage Analyzer Pro",
         page_icon="🏗️",
         layout="wide",
         initial_sidebar_state="expanded"
     )
 
-    # Custom CSS
-    st.markdown("""
-        <style>
-        .main {
-            padding: 2rem;
-        }
-        .stProgress > div > div > div > div {
-            background-image: linear-gradient(to right, #ff6b6b, #f06595);
-        }
-        .damage-card {
-            padding: 1.5rem;
-            border-radius: 0.5rem;
-            background: #f8f9fa;
-            margin-bottom: 1rem;
-        }
-        </style>
-    """, unsafe_allow_html=True)
+    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()
 
-    col1, col2, col3 = st.columns([1, 3, 1])
-    with col2:
-        st.title("🏗️ Structural Damage Analyzer")
-        st.markdown("### Upload a photo of structural damage for instant analysis")
+    # Load model
+    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()
-    
-    upload_col1, upload_col2, upload_col3 = st.columns([1, 2, 1])
-    with upload_col2:
-        uploaded_file = st.file_uploader("Choose an image file", type=['jpg', 'jpeg', 'png'])
+    # File upload with drag and drop
+    uploaded_file = st.file_uploader(
+        "Drag and drop or click to upload an image",
+        type=['jpg', 'jpeg', 'png'],
+        help="Supported formats: JPG, JPEG, PNG"
+    )
 
     if uploaded_file:
-        image = Image.open(uploaded_file)
-        
-        analysis_col1, analysis_col2 = st.columns(2)
-        
-        with analysis_col1:
-            st.image(image, caption="Uploaded Structure", use_column_width=True)
-        
-        with analysis_col2:
-            with st.spinner("🔍 Analyzing structural damage..."):
-                predictions = analyze_damage(image, model, processor)
-                
-                st.markdown("### 📊 Damage Assessment Results")
-                
-                for idx, prob in enumerate(predictions):
-                    confidence = float(prob) * 100
-                    if confidence > 15:
-                        damage_type = DAMAGE_TYPES[idx]['name']
-                        cases = KNOWLEDGE_BASE[damage_type]
-                        
-                        st.markdown(f"""
-                        <div class="damage-card">
-                            <h4>{damage_type.replace('_', ' ').title()} Detected</h4>
-                        </div>
-                        """, unsafe_allow_html=True)
+        try:
+            # Image display and analysis
+            image = Image.open(uploaded_file)
+            col1, col2 = st.columns([1, 1])
+            
+            with col1:
+                st.image(image, caption="Uploaded Structure", use_column_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:
+                        st.markdown("### 📊 Analysis Results")
+                        st.markdown(f"*Analysis completed in {analysis_time:.2f} seconds*")
                         
-                        col1, col2 = st.columns([3, 1])
-                        with col1:
-                            st.progress(confidence / 100)
-                        with col2:
-                            st.write(f"Confidence: {confidence:.1f}%")
+                        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):
+                                    # Progress bar with custom color
+                                    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'])
                         
-                        tabs = st.tabs(["Details", "Repair Methods", "Recommendations"])
+                        if not detected:
+                            st.info("No significant structural damage detected. Regular maintenance recommended.")
                         
-                        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]:
-                            st.markdown("#### Repair Steps")
-                            for step in cases[0]['repair_method']:
-                                st.markdown(f"- {step}")
-                            st.markdown(f"""
-                            - **Estimated Cost:** {cases[0]['estimated_cost']}
-                            - **Timeframe:** {cases[0]['timeframe']}
-                            """)
-                        
-                        with tabs[2]:
-                            st.markdown("#### Immediate Actions")
-                            st.warning(cases[0]['immediate_action'])
-                            
-                            st.markdown("#### Prevention Measures")
-                            st.info(cases[0]['prevention'])
+                        # Add to history
+                        st.session_state.history.append(f"Analyzed image: {uploaded_file.name}")
+
+        except Exception as e:
+            st.error(f"Error processing image: {str(e)}")
+            st.info("Please try uploading a different image.")
 
     # Footer
     st.markdown("---")
-    st.markdown("""
+    st.markdown(
+        """
         <div style='text-align: center'>
-            <p>🏗️ Structural Damage Analysis Tool | Built with Streamlit</p>
+            <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()