import streamlit as st
import torch
import cv2
import numpy as np
from PIL import Image
from groq import Groq
import os

# Load Pretrained Model (YOLOv8 for Structural Damage Detection)
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

def process_image(image):
    # Convert PIL image to OpenCV format
    img = np.array(image)
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    
    # Perform inference
    results = model(img)
    
    # Draw bounding boxes on detected damages
    for *box, conf, cls in results.xyxy[0].tolist():
        x1, y1, x2, y2 = map(int, box)
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
        label = f"{model.names[int(cls)]}: {conf:.2f}"
        cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    return img

# Initialize Groq API Client
groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))

def get_ai_analysis():
    chat_completion = groq_client.chat.completions.create(
        messages=[
            {"role": "user", "content": "Explain the impact of structural damage in buildings."}
        ],
        model="llama-3.3-70b-versatile",
        stream=False,
    )
    return chat_completion.choices[0].message.content

# Streamlit UI
st.title("AI-Based Structural Damage Detection")
st.write("Upload an image of a structure, and the AI will detect any visible structural damage.")

uploaded_file = st.file_uploader("Choose an image", type=["jpg", "png", "jpeg"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_container_width=True)
    
    if st.button("Analyze Image"):
        processed_img = process_image(image)
        st.image(processed_img, caption="Detected Structural Damage", use_container_width=True)
        
        st.subheader("AI Analysis")
        st.write(get_ai_analysis())