File size: 2,715 Bytes
15526fe 4f207d0 15526fe 9814e5f 15526fe 4f207d0 44bef16 15526fe 4f207d0 15526fe 9814e5f 44bef16 15526fe 4f207d0 44bef16 15526fe 9814e5f 4f207d0 15526fe 4f207d0 44bef16 4f207d0 9814e5f |
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 |
import streamlit as st
from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
import torch
# 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'
}
}
@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="Analyzed Structure")
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']}")
st.write("Recommended Actions:")
for remedy in DAMAGE_TYPES[damage_type]['remedies']:
st.write(f"• {remedy}")
st.divider()
if __name__ == "__main__":
main() |