Update app.py

app.py

@@ -1,46 +1,71 @@
 import streamlit as st
-from transformers import pipeline, AutoModelForImageClassification, AutoImageProcessor
+from transformers import ViTForImageClassification, ViTImageProcessor
 from PIL import Image
 import torch
 
-# Initialize models
-model_name = "microsoft/resnet-50"
-image_processor = AutoImageProcessor.from_pretrained(model_name)
-model = AutoModelForImageClassification.from_pretrained(model_name)
-classifier = pipeline("image-classification", model=model, image_processor=image_processor)
-
-damage_types = {
-    'spalling': ['Remove loose material', 'Apply repair mortar', 'Seal surface'],
-    'cracking': ['Measure crack width', 'Install crack monitors', 'Apply sealant'],
-    'corrosion': ['Clean exposed steel', 'Apply rust converter', 'Paint protective coating'],
-    'dampness': ['Locate water source', 'Fix drainage issues', 'Apply waterproofing']
+# 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'
+    }
 }
 
-def analyze_image(image):
-    return classifier(image)
+@st.cache_resource
+def load_model():
+    model = ViTForImageClassification.from_pretrained(
+        "nateraw/vit-base-patch16-224-finetuned-structural-damage",
+        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):
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+    return probs[0]
 
 def main():
     st.title("Structural Damage Assessment")
     
+    model, processor = load_model()
+    
     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 Image")
+        st.image(image, caption="Analyzed Structure")
         
-        with st.spinner("Analyzing..."):
-            results = analyze_image(image)
+        with st.spinner("Analyzing structural damage..."):
+            predictions = analyze_damage(image, model, processor)
             
-            for result in results:
-                confidence = result['score'] * 100
-                st.progress(confidence / 100)
-                st.write(f"{result['label']}: {confidence:.1f}%")
-                
-                for damage_type, remedies in damage_types.items():
-                    if damage_type in result['label'].lower():
-                        st.subheader("Recommended Actions:")
-                        for remedy in remedies:
-                            st.write(f"• {remedy}")
+            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']}")
+                    
+                    st.write("Recommended Actions:")
+                    for remedy in DAMAGE_TYPES[damage_type]['remedies']:
+                        st.write(f"• {remedy}")
+                    st.divider()
 
 if __name__ == "__main__":
     main()