import streamlit as st
import kagglehub
import tensorflow as tf
from PIL import Image
import numpy as np

# Download model from Kaggle Hub
model_path = kagglehub.model_download("sohiebalwedyan/seatbelt-detection/tensorFlow2/seatbelt-detection")
print("Path to model files:", model_path)

# Load the TensorFlow model
model = tf.saved_model.load(model_path)

def predict(image):
    # Preprocess image
    img = image.resize((224, 224))
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    
    # Run inference
    predictions = model(tf.convert_to_tensor(img_array, dtype=tf.float32))
    
    # Process predictions (modify based on actual model output)
    # This is a placeholder and may need adjustment based on model specifics
    st.write("Raw predictions:", predictions)
    
    return image  # Placeholder return

# Streamlit UI
st.title("Seatbelt Detection")
st.markdown("Upload an image for seatbelt detection.")

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

if uploaded_image is not None:
    image = Image.open(uploaded_image)
    st.image(image, caption="Uploaded Image", use_column_width=True)
    
    st.subheader("Prediction Results:")
    result = predict(image)
    # Note: You may need to modify visualization based on model output