Update app.py

app.py

@@ -2,70 +2,104 @@ import streamlit as st
 from transformers import ViTForImageClassification, ViTImageProcessor
 from PIL import Image
 import torch
+import numpy as np
 
-# Custom labels for structural damage
 DAMAGE_TYPES = {
-    'spalling': {
-        'remedies': ['Remove loose concrete', 'Clean exposed reinforcement', 'Apply repair mortar'],
-        'risk_level': 'High'
-    },
-    'reinforcement_corrosion': {
-        'remedies': ['Remove rust', 'Apply corrosion inhibitor', 'Repair concrete cover'],
-        'risk_level': 'Critical'
-    },
-    'structural_crack': {
-        'remedies': ['Monitor crack width', 'Epoxy injection', 'Structural reinforcement'],
-        'risk_level': 'Critical'
-    },
-    'dampness': {
-        'remedies': ['Identify water source', 'Fix waterproofing', 'Improve drainage'],
-        'risk_level': 'Medium'
-    }
+    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'}
+}
+
+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'
+    ]
 }
 
 @st.cache_resource
 def load_model():
     model = ViTForImageClassification.from_pretrained(
-        "nateraw/vit-base-patch16-224-finetuned-structural-damage",
+        "google/vit-base-patch16-224",
         num_labels=len(DAMAGE_TYPES),
-        id2label={i: label for i, label in enumerate(DAMAGE_TYPES.keys())},
-        label2id={label: i for i, label in enumerate(DAMAGE_TYPES.keys())}
     )
     processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
     return model, processor
 
-def analyze_damage(image, model, processor):
+def process_image(image, 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)
-    return probs[0]
+    probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
+    return probs
 
 def main():
-    st.title("Structural Damage Assessment")
-    
+    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("Upload structural image", type=['jpg', 'jpeg', 'png'])
+    uploaded_file = st.file_uploader("Choose an image", type=['jpg', 'jpeg', 'png'])
     
     if uploaded_file:
         image = Image.open(uploaded_file)
-        st.image(image, caption="Analyzed Structure")
+        st.image(image, caption="Uploaded Structure", use_column_width=True)
         
         with st.spinner("Analyzing structural damage..."):
             predictions = analyze_damage(image, model, processor)
             
-            for idx, (damage_type, conf) in enumerate(zip(DAMAGE_TYPES.keys(), predictions)):
-                confidence = float(conf) * 100
-                if confidence > 10:  # Show only significant predictions
-                    st.subheader(f"{damage_type.replace('_', ' ').title()}")
-                    st.progress(confidence / 100)
-                    st.write(f"Confidence: {confidence:.1f}%")
-                    st.write(f"Risk Level: {DAMAGE_TYPES[damage_type]['risk_level']}")
+            col1, col2 = st.columns(2)
+            
+            with col1:
+                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}")
+            
+            with col2:
+                st.subheader("Recommended Actions")
+                for idx, prob in enumerate(predictions):
+                    damage_type = DAMAGE_TYPES[idx]['name']
+                    confidence = float(prob) * 100
                     
-                    st.write("Recommended Actions:")
-                    for remedy in DAMAGE_TYPES[damage_type]['remedies']:
-                        st.write(f"• {remedy}")
-                    st.divider()
+                    if confidence > 15:
+                        st.write(f"**For {damage_type.replace('_', ' ').title()}:**")
+                        for remedy in REMEDIES[damage_type]:
+                            st.write(f"• {remedy}")
+                        st.write("---")
 
 if __name__ == "__main__":
     main()