File size: 11,428 Bytes
3db0aec
15526fe
3db0aec
 
9814e5f
3db0aec
 
 
 
15526fe
3db0aec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dbd2162
60dab82
3db0aec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
15526fe
3db0aec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dbd2162
3db0aec
dbd2162
3db0aec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dbd2162
3db0aec
 
 
 
 
dbd2162
3db0aec
 
 
 
 
 
 
 
 
4f207d0
3db0aec
 
 
 
dbd2162
3db0aec
 
 
 
 
 
 
dbd2162
3db0aec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dbd2162
3db0aec
 
15526fe
9814e5f
3db0aec
 
 
 
 
 
 
 
 
 
dbd2162
3db0aec
 
 
 
 
 
 
 
dbd2162
3db0aec
 
 
 
 
 
 
 
 
 
4f207d0
3db0aec
 
 
 
 
dbd2162
3db0aec
dbd2162
3db0aec
 
 
 
 
 
dbd2162
3db0aec
 
 
dbd2162
3db0aec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9814e5f
 
dbd2162
3db0aec
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
```python
import streamlit as st
from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
import torch
import time
import gc
from knowledge_base import KNOWLEDGE_BASE, DAMAGE_TYPES
from rag_utils import RAGSystem

# Constants
MAX_FILE_SIZE = 5 * 1024 * 1024  # 5MB
MAX_IMAGE_SIZE = 1024  # Maximum dimension for images

# Cache the model and RAG system globally
MODEL = None
PROCESSOR = None
RAG_SYSTEM = None

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

@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"
        model = ViTForImageClassification.from_pretrained(
            model_name,
            num_labels=len(DAMAGE_TYPES),
            ignore_mismatched_sizes=True,
            device_map="auto"
        )
        processor = ViTImageProcessor.from_pretrained(model_name)
        return model, processor
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None, None

def init_rag_system():
    """Initialize the RAG system with knowledge base"""
    global RAG_SYSTEM
    if RAG_SYSTEM is None:
        RAG_SYSTEM = RAGSystem()
        RAG_SYSTEM.initialize_knowledge_base(KNOWLEDGE_BASE)

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):
    """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)}")
        return None

def analyze_damage(image, model, processor):
    """Analyze structural damage in the image"""
    try:
        with torch.no_grad():
            image = image.convert('RGB')
            inputs = processor(images=image, return_tensors="pt")
            outputs = model(**inputs)
            probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
            cleanup_memory()
            return probs
    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 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
    )

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?"
        )
        
        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.")

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"
    )

    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()

    # 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
    
    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"
    )

    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}")

        except Exception as e:
            cleanup_memory()
            st.error(f"Error processing image: {str(e)}")
            st.info("Please try uploading a different 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
    )

if __name__ == "__main__":
    main()
```