import streamlit as st
import base64
import openai
import numpy as np
import cv2
from tensorflow.keras.models import load_model
from keras.preprocessing.image import img_to_array
from keras.applications.inception_v3 import preprocess_input
import os
from PIL import Image
import io
st.set_page_config(page_title="Wall Defect Classifier", layout="centered")

# Set OpenAI API Key
openai.api_key = os.getenv('OPEN_AI')

# Defect categories
class_labels = [
    "Floor Cracks",
    "Floor Peel",
    "Shade Fading",
    "Air Bubble (Single Defect)",
    "Spike Roller (Single Defect)"
]

@st.cache_resource
def load_trained_model():
    return load_model('my_inceptionmodelwithoutaug (1).h5')

def compress_image(image_bytes, max_size_kb=500):
    img = Image.open(io.BytesIO(image_bytes))
    quality = 95
    output_bytes = io.BytesIO()
    while True:
        output_bytes.seek(0)
        output_bytes.truncate()
        img.save(output_bytes, format='JPEG', quality=quality)
        if len(output_bytes.getvalue()) <= max_size_kb * 1024 or quality <= 5:
            break
        quality -= 5
    return output_bytes.getvalue()

def process_image_for_openai(image_bytes):
    compressed_image = compress_image(image_bytes)
    return base64.b64encode(compressed_image).decode('utf-8')

# Load model once
loaded_model = load_trained_model()

# App UI



st.title("🧠 Wall Defect Classification & AI-Based Description")
category_choice = st.selectbox("🛠️ Select Defect Category Type:", ["Flooring"], index=0)

st.markdown("Upload a wall surface image to detect potential defects and generate a structured AI analysis.")


uploaded_file = st.file_uploader("📤 Upload an Image", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    file_bytes = uploaded_file.getvalue()

    st.image(file_bytes, caption="🖼️ Uploaded Image", use_column_width=True)

    # Preprocess for prediction
    input_img = cv2.imdecode(np.frombuffer(file_bytes, np.uint8), cv2.IMREAD_COLOR)
    input_img_resized = cv2.resize(input_img, (256, 256))
    x = img_to_array(input_img_resized)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)

    preds = loaded_model.predict(x)
    class_index = np.argmax(preds[0])
    max_probability = preds[0][class_index]
    class_name = class_labels[class_index]

    # Classification Result Display
    st.subheader("🔍 Classification Result")
    st.success(f"**Predicted Defect:** {class_name}")
    st.progress(min(int(max_probability * 100), 100))
    st.markdown(f"**Confidence Level:** {max_probability:.2%}")

    if max_probability < 0.59:
        st.warning("⚠️ The confidence score is below 59%. Please manually verify this result.")
    else:
        compressed_base64 = process_image_for_openai(file_bytes)
        ai_prompt = (
             f"Our trained model predicts the following defect: {class_name}. "
            f"Can you analyze the following image and generate AI-based descriptions "
           
            f"for this defect? The output format should be:\n"
            f"Category ID: <Insert Category ID>\n"
            f"Title: <Short Title of the Defect>\n"
            f"Description: <A concise, technical description in 100 words or less>"
        )
    
        st.subheader("Generating AI-Based Analysis...")
        try:
            response = openai.ChatCompletion.create(
                model="gpt-4o",
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": ai_prompt},
                            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{compressed_base64}"}}
                        ]
                    }
                ],
                max_tokens=300,
            )
            ai_description = response.choices[0].message.content
            st.subheader("📝 AI-Generated Defect Description")
            st.text_area("Output", value=ai_description.strip(), height=250)
        except Exception as e:
          st.error(f"❌ An error occurred while generating the description:\n{e}")