Keshav-rejoice's picture
Update app.py
2ddb5e6 verified
raw
history blame
4.14 kB
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}")