app.py

import streamlit as st
from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
import torch
import time
import gc

# Constants
MAX_FILE_SIZE = 5 * 1024 * 1024  # 5MB
MAX_IMAGE_SIZE = 1024  # Maximum dimension for images

def cleanup_memory():
    """Clean up memory and GPU cache"""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

def init_session_state():
    """Initialize session state variables"""
    if 'history' not in st.session_state:
        st.session_state.history = []
    if 'dark_mode' not in st.session_state:
        st.session_state.dark_mode = False

@st.cache_resource(show_spinner="Loading AI model...")
def load_model():
    """Load and cache the model and processor"""
    try:
        model_name = "google/vit-base-patch16-224"
        processor = ViTImageProcessor.from_pretrained(model_name)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model = ViTForImageClassification.from_pretrained(
            model_name,
            num_labels=len(DAMAGE_TYPES),
            ignore_mismatched_sizes=True,
        ).to(device)
        model.eval()
        return model, processor
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None, None

def validate_image(image):
    """Validate image size and format"""
    if image.size[0] * image.size[1] > 1024 * 1024:
        st.warning("Large image detected. The image will be resized for better performance.")
    if image.format not in ['JPEG', 'PNG']:
        st.warning("Image format not optimal. Consider using JPEG or PNG for better performance.")

def preprocess_image(uploaded_file):
    """Preprocess and validate uploaded image"""
    try:
        image = Image.open(uploaded_file)
        if max(image.size) > MAX_IMAGE_SIZE:
            ratio = MAX_IMAGE_SIZE / max(image.size)
            new_size = tuple([int(dim * ratio) for dim in image.size])
            image = image.resize(new_size, Image.Resampling.LANCZOS)
        return image
    except Exception as e:
        st.error(f"Error processing image: {str(e)}")
        return None

def analyze_damage(image, model, processor):
    """Analyze structural damage in the image"""
    try:
        device = next(model.parameters()).device
        with torch.no_grad():
            image = image.convert('RGB')
            inputs = processor(images=image, return_tensors="pt")
            inputs = {k: v.to(device) for k, v in inputs.items()}
            outputs = model(**inputs)
            probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
            cleanup_memory()
            return probs.cpu()
    except RuntimeError as e:
        if "out of memory" in str(e):
            cleanup_memory()
            st.error("Out of memory. Please try with a smaller image.")
        else:
            st.error(f"Error analyzing image: {str(e)}")
        return None

def main():
    st.set_page_config(
        page_title="Structural Damage Analyzer Pro",
        page_icon="🏗️",
        layout="wide",
        initial_sidebar_state="expanded"
    )

    # Initialize session state
    init_session_state()

    # Display header
    st.markdown(
        """
        <div style='text-align: center; padding: 1rem;'>
            <h1>🏗️ Structural Damage Analyzer Pro</h1>
            <p style='font-size: 1.2rem;'>Advanced AI-powered structural damage assessment tool</p>
        </div>
        """,
        unsafe_allow_html=True
    )

    # Load model
    model, processor = load_model()
    if model is None:
        st.error("Failed to load model. Please refresh the page.")
        return

    # File upload
    uploaded_file = st.file_uploader(
        "Upload an image for analysis",
        type=['jpg', 'jpeg', 'png'],
        help="Supported formats: JPG, JPEG, PNG"
    )

    if uploaded_file:
        try:
            if uploaded_file.size > MAX_FILE_SIZE:
                st.error("File size too large. Please upload an image smaller than 5MB.")
                return

            image = preprocess_image(uploaded_file)
            if image is None:
                return

            validate_image(image)
            
            # Display image and analyze
            st.image(image, caption="Uploaded Structure", use_column_width=True)
            
            with st.spinner("🔍 Analyzing damage..."):
                predictions = analyze_damage(image, model, processor)
                if predictions is not None:
                    st.success("Analysis complete!")
                    # Add analysis display logic here based on your DAMAGE_TYPES

        except Exception as e:
            cleanup_memory()
            st.error(f"Error processing image: {str(e)}")
            st.info("Please try uploading a different image.")

    # Footer
    st.markdown("---")
    st.markdown(
        """
        <div style='text-align: center'>
            <p>🏗️ Structural Damage Analyzer Pro | Built with Streamlit & Transformers</p>
            <p style='font-size: 0.8rem;'>For professional use only. Always consult with a structural engineer.</p>
        </div>
        """,
        unsafe_allow_html=True
    )

if __name__ == "__main__":
    main()