Update app.py

app.py

@@ -2,8 +2,8 @@ import streamlit as st
 from transformers import ViTForImageClassification, ViTImageProcessor
 from PIL import Image
 import torch
-import numpy as np
 
+# Define damage types and remedies
 DAMAGE_TYPES = {
     0: {'name': 'spalling', 'risk': 'High'},
     1: {'name': 'reinforcement_corrosion', 'risk': 'Critical'},
@@ -13,30 +13,11 @@ DAMAGE_TYPES = {
 }
 
 REMEDIES = {
-    'spalling': [
-        'Remove loose concrete',
-        'Clean exposed area',
-        'Apply repair mortar'
-    ],
-    'reinforcement_corrosion': [
-        'Remove rust',
-        'Apply corrosion inhibitor',
-        'Repair concrete cover'
-    ],
-    'structural_crack': [
-        'Measure crack width',
-        'Epoxy injection',
-        'Monitor progression'
-    ],
-    'dampness': [
-        'Identify water source',
-        'Improve drainage',
-        'Apply waterproofing'
-    ],
-    'no_damage': [
-        'Regular maintenance',
-        'Periodic inspection'
-    ]
+    'spalling': ['Remove loose concrete', 'Clean exposed area', 'Apply repair mortar'],
+    'reinforcement_corrosion': ['Remove rust', 'Apply corrosion inhibitor', 'Repair concrete cover'],
+    'structural_crack': ['Measure crack width', 'Epoxy injection', 'Monitor progression'],
+    'dampness': ['Identify water source', 'Improve drainage', 'Apply waterproofing'],
+    'no_damage': ['Regular maintenance', 'Periodic inspection']
 }
 
 @st.cache_resource
@@ -44,34 +25,30 @@ 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
 
-def process_image(image, processor):
+def analyze_damage(image, model, processor):
     image = image.convert('RGB')
     inputs = processor(images=image, return_tensors="pt")
-    return inputs
-
-def analyze_damage(image, model, processor):
-    inputs = process_image(image, processor)
     outputs = model(**inputs)
     probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
     return probs
 
 def main():
     st.title("Structural Damage Assessment Tool")
-    st.write("Upload an image of building structure for damage analysis")
 
     model, processor = load_model()
     
-    uploaded_file = st.file_uploader("Choose an image", type=['jpg', 'jpeg', 'png'])
+    uploaded_file = st.file_uploader("Upload structural image", type=['jpg', 'jpeg', 'png'])
     
     if uploaded_file:
         image = Image.open(uploaded_file)
         st.image(image, caption="Uploaded Structure", use_column_width=True)
         
-        with st.spinner("Analyzing structural damage..."):
+        with st.spinner("Analyzing..."):
             predictions = analyze_damage(image, model, processor)
             
             col1, col2 = st.columns(2)
@@ -80,21 +57,18 @@ def main():
                 st.subheader("Damage Assessment")
                 for idx, prob in enumerate(predictions):
                     damage_type = DAMAGE_TYPES[idx]['name']
-                    risk_level = DAMAGE_TYPES[idx]['risk']
                     confidence = float(prob) * 100
-                    
                     if confidence > 15:
                         st.write(f"**{damage_type.replace('_', ' ').title()}**")
                         st.progress(confidence / 100)
                         st.write(f"Confidence: {confidence:.1f}%")
-                        st.write(f"Risk Level: {risk_level}")
+                        st.write(f"Risk Level: {DAMAGE_TYPES[idx]['risk']}")
             
             with col2:
                 st.subheader("Recommended Actions")
                 for idx, prob in enumerate(predictions):
                     damage_type = DAMAGE_TYPES[idx]['name']
                     confidence = float(prob) * 100
-                    
                     if confidence > 15:
                         st.write(f"**For {damage_type.replace('_', ' ').title()}:**")
                         for remedy in REMEDIES[damage_type]: