import streamlit as st
import cv2
import numpy as np
from tensorflow.keras.models import load_model

# Model loading
model = load_model('model.h5')

# Specifying Height and Width
IMAGE_HEIGHT, IMAGE_WIDTH = 128, 128

SEQUENCE_LENGTH = 12

CLASSES_LIST = ["NonViolence", "Violence"]

def predict_video(video_file_path, sequence_length):

    video_reader = cv2.VideoCapture(video_file_path)

    # Get original width and height of the video.
    original_video_width = int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH))
    original_video_height = int(video_reader.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # List to store video frames
    frames_list = []

    # Storing the predicted class
    predicted_class_name = ''

    # Get total number of frames
    video_frames_count = int(video_reader.get(cv2.CAP_PROP_FRAME_COUNT))

    # interval after which frames will be added
    skip_frames_window = max(int(video_frames_count / sequence_length), 1)

    for frame_counter in range(sequence_length):
        # set the current frame
        video_reader.set(cv2.CAP_PROP_POS_FRAMES, frame_counter * skip_frames_window)

        success, frame = video_reader.read()

        if not success:
            break

        # resize the Frame
        resized_frame = cv2.resize(frame, (IMAGE_HEIGHT, IMAGE_WIDTH))

        # normalizing the resized frame.
        normalized_frame = resized_frame / 255

        # appending frames into the frames list
        frames_list.append(normalized_frame)

    # passing frames to model and get predictions
    predicted_labels_probabilities = model.predict(np.expand_dims(frames_list, axis=0))[0]

    # get index of class with highest probability.
    predicted_label = np.argmax(predicted_labels_probabilities)

    # Get the class name using the retrieved index.
    predicted_class_name = CLASSES_LIST[predicted_label]

    video_reader.release()

    return predicted_class_name, predicted_labels_probabilities[predicted_label]


def main():
    st.title("Violence Detection")

    extensions = ['mp4', 'png', 'avi', 'mov']
    # Upload a video file
    uploaded_file = st.file_uploader("Upload a video file", type=extensions)

    if uploaded_file:
        # Get video file path
        video_file_path = "temp_video.mp4"
        with open(video_file_path, "wb") as f:
            f.write(uploaded_file.getvalue())

        # Predict on the video
        predicted_class, confidence = predict_video(video_file_path, SEQUENCE_LENGTH)

        # Display the predicted class along with the prediction confidence.
        st.write(f'Predicted: {predicted_class}\nConfidence: {confidence}')

        # Play the actual video
        st.video(video_file_path)


if __name__ == "__main__":
    main()