Spaces:
Sleeping
Sleeping
| import os | |
| os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0' | |
| import tensorflow as tf | |
| import cv2 | |
| import numpy as np | |
| import math | |
| import gradio as gr | |
| class ShopliftingPrediction: | |
| def __init__(self, model_path, frame_width, frame_height, sequence_length): | |
| self.frame_width = frame_width | |
| self.frame_height = frame_height | |
| self.sequence_length = sequence_length | |
| self.model_path = model_path | |
| self.message = '' | |
| def load_model(self): | |
| custom_objects = { | |
| 'Conv2D': tf.keras.layers.Conv2D, | |
| 'MaxPooling2D': tf.keras.layers.MaxPooling2D, | |
| 'TimeDistributed': tf.keras.layers.TimeDistributed, | |
| 'LSTM': tf.keras.layers.LSTM, | |
| 'Dense': tf.keras.layers.Dense, | |
| 'Flatten': tf.keras.layers.Flatten, | |
| 'Dropout': tf.keras.layers.Dropout, | |
| 'Orthogonal': tf.keras.initializers.Orthogonal, | |
| } | |
| self.model = tf.keras.models.load_model(self.model_path, custom_objects=custom_objects) | |
| def generate_message_content(self, probability, label): | |
| if label == 0: | |
| if probability <= 50: | |
| self.message = "No theft" | |
| elif probability <= 75: | |
| self.message = "There is little chance of theft" | |
| elif probability <= 85: | |
| self.message = "High probability of theft" | |
| else: | |
| self.message = "Very high probability of theft" | |
| elif label == 1: | |
| if probability <= 50: | |
| self.message = "No theft" | |
| elif probability <= 75: | |
| self.message = "The movement is confusing, watch" | |
| elif probability <= 85: | |
| self.message = "I think it's normal, but it's better to watch" | |
| else: | |
| self.message = "Movement is normal" | |
| def Pre_Process_Video(self, current_frame, previous_frame): | |
| diff = cv2.absdiff(current_frame, previous_frame) | |
| diff = cv2.GaussianBlur(diff, (3, 3), 0) | |
| resized_frame = cv2.resize(diff, (self.frame_height, self.frame_width)) | |
| gray_frame = cv2.cvtColor(resized_frame, cv2.COLOR_BGR2GRAY) | |
| normalized_frame = gray_frame / 255 | |
| return normalized_frame | |
| def Single_Frame_Predict(self, frames_queue): | |
| probabilities = self.model.predict(np.expand_dims(frames_queue, axis=0))[0] | |
| predicted_label = np.argmax(probabilities) | |
| probability = math.floor(max(probabilities[0], probabilities[1]) * 100) | |
| return [probability, predicted_label] | |
| def process_frame(self, frame, frames_queue, previous): | |
| normalized_frame = self.Pre_Process_Video(frame, previous) | |
| frames_queue.append(normalized_frame) | |
| if len(frames_queue) == self.sequence_length: | |
| [probability, predicted_label] = self.Single_Frame_Predict(frames_queue) | |
| self.generate_message_content(probability, predicted_label) | |
| frames_queue = [] | |
| cv2.rectangle(frame, (0, 0), (640, 40), (255, 255, 255), -1) | |
| cv2.putText(frame, self.message, (1, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA) | |
| return frame | |
| def inference(model_path): | |
| shoplifting_prediction = ShopliftingPrediction(model_path, 64, 64, 30) | |
| shoplifting_prediction.load_model() | |
| def process_video(video_frame): | |
| frame = shoplifting_prediction.process_frame(video_frame, [], video_frame) | |
| return frame | |
| return process_video | |
| model_path = 'lrcn_160S_90_90Q.h5' | |
| process_video = inference(model_path) | |
| iface = gr.Interface( | |
| fn=process_video, | |
| inputs=gr.inputs.Video(type="numpy"), | |
| outputs="video", | |
| live=True, | |
| capture_session=True # Ensures the video stream is captured properly | |
| ) | |
| iface.launch() | |