|
|
import streamlit as st |
|
|
from transformers import ViTForImageClassification, ViTImageProcessor |
|
|
from PIL import Image |
|
|
import torch |
|
|
|
|
|
|
|
|
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: |
|
|
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() |
