import streamlit as st
import numpy as np
import torch
import cv2
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image as keras_image
from PIL import Image

# Load the pre-trained MobileNetV2 model from the .h5 file
@st.cache_resource
def load_mobilenet_model():
    return load_model('mobilenetv2_classifier.h5')  # Replace with your .h5 file path

mobilenet_model = load_mobilenet_model()

# Load the YOLOv5 model (pre-trained)
@st.cache_resource
def load_yolov5_model():
    return torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

yolov5_model = load_yolov5_model()

# Function to preprocess and predict image using MobileNetV2 with class messages
def predict_image(model, img, target_size=(224, 224)):
    st.image(img, caption="Uploaded Image", use_column_width=True)
    img = img.resize(target_size)
    img_array = keras_image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array /= 255.0  # Normalize to match the training preprocessing
    predictions = model.predict(img_array)[0]

    # Get indices sorted by prediction confidence (in descending order)
    sorted_indices = np.argsort(predictions)[::-1]
    primary_class_index = sorted_indices[0]
    primary_confidence = predictions[primary_class_index]

    # Check if the primary predicted label is 1 and has a confidence below 0.7
    if primary_class_index != 1 and primary_confidence < 0.7:
        secondary_class_index = sorted_indices[1]
        secondary_confidence = predictions[secondary_class_index]
        return secondary_class_index, secondary_confidence

    return primary_class_index, primary_confidence

# Map class indices to specific messages
class_messages = {
    0: "Face is not visible",
    1: "Looking at the system",
    2: "Looking down",
    3: "Looking left or right",
    4: "Looking up"
}

# Function to count people using YOLOv5
def count_people(img):
    img_array = np.array(img)
    results = yolov5_model(img_array)
    detections = results.pandas().xyxy[0]
    person_detections = detections[detections['name'] == 'person']
    num_people = len(person_detections)
    return num_people, results

# Function to detect communication devices using YOLOv5
def detect_device(img):
    img_array = np.array(img)
    results = yolov5_model(img_array)
    detections = results.pandas().xyxy[0]
    device_detections = detections[detections['name'].str.contains('phone|tablet|computer', case=False, na=False)]
    return device_detections, results

# Streamlit UI setup
st.title("Image Analysis App")

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

if uploaded_file is not None:
    img = Image.open(uploaded_file)
    
    # Display the uploaded image
    st.image(img, caption="Uploaded Image", use_column_width=True)
    
    # Run classification with MobileNetV2
    st.subheader("Image Classification")
    predicted_index, confidence = predict_image(mobilenet_model, img)
    
    if predicted_index in class_messages:
        st.write(f"Prediction: {class_messages[predicted_index]}")
        st.write(f"Confidence score: {confidence:.2f}")
    else:
        st.write("Prediction not recognized.")
    
    # Run person counting with YOLOv5
    st.subheader("People Counting")
    num_people, people_results = count_people(img)
    st.write(f"Number of people detected: {num_people}")
    people_results.show()  # Show image with bounding boxes

    # Run device detection with YOLOv5
    st.subheader("Communication Device Detection")
    device_detections, device_results = detect_device(img)
    
    if not device_detections.empty:
        st.write("Communication devices detected:")
        st.write(device_detections[['name', 'confidence']])
        device_results.show()  # Show image with bounding boxes
    else:
        st.write("No communication devices detected in the image.")
else:
    st.write("Please upload an image to analyze.")