import streamlit as st
from PIL import Image
import cv2
import numpy as np
from ultralytics import YOLO
import os

# Model Initialization

model = YOLO(https://colab.research.google.com/drive/1IAyIjPN1J_9s5MhJ0uCsccxfcA0WfXl2?authuser=1#scrollTo=bA29p0kfmhlP)

# Adjust the path to your YOLOv11 model
# Function to detect actions in images
def detect_action(image_path):
    results = model.predict(source=image_path, conf=0.25, save=False)
    result = results[0]
    detections = [
        (model.names[int(box.cls[0])], float(box.conf[0])) for box in result.boxes
    ]
    
    # Classify action based on detections
    action_scores = classify_action(detections)
    
    return result.plot(), action_scores

def classify_action(detections):
    detected_objects = [d[0] for d in detections]
    
    action_scores = {
        'Stealing': 0.0,
        'Sneaking': 0.0,
        'Peaking': 0.0,
        'Normal': 0.0
    }

    if 'person' in detected_objects:
        if any(obj in detected_objects for obj in ['backpack', 'handbag', 'suitcase']):
            action_scores['Stealing'] += 0.4
        if 'refrigerator' in detected_objects:
            action_scores['Stealing'] += 0.3
        if [conf for obj, conf in detections if obj == 'person'][0] < 0.6:
            action_scores['Sneaking'] += 0.5
        if len(detected_objects) <= 2:
            action_scores['Peaking'] += 0.5

    if not any(score > 0.3 for score in action_scores.values()):
        action_scores['Normal'] = 0.4

    return action_scores

# Streamlit UI
st.title('Suspicious Activity Detection')
st.write('Upload an image to detect suspicious activities.')

# File uploader
uploaded_file = st.file_uploader("Choose an image...", type="jpg")
if uploaded_file is not None:
    # Read the image
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image', use_column_width=True)
    
    # Save the uploaded file for processing
    img_path = "/tmp/uploaded_image.jpg"
    image.save(img_path)
    
    # Predict and display results
    st.write("Detecting action...")
    detected_image, action_scores = detect_action(img_path)
    
    st.image(detected_image, caption='Detected Image', use_column_width=True)
    
    # Display action scores
    st.write("Action Probability Scores:")
    for action, score in action_scores.items():
        st.write(f"{action}: {score:.2%}")
    
    # Predict and display the most likely action
    predicted_action = max(action_scores.items(), key=lambda x: x[1])
    st.write(f"Predicted Action: {predicted_action[0]} ({predicted_action[1]:.2%} confidence)")