Abs6187's picture
Update app.py
fddfd57 verified
raw
history blame
2.51 kB
import gradio as gr
import cv2
import numpy as np
from tensorflow.keras.models import load_model
from sklearn.preprocessing import StandardScaler
from ultralytics import YOLO
# Load models
lstm_model = load_model('suspicious_activity_model.h5')
yolo_model = YOLO('yolov8n-pose.pt') # Ensure this model supports keypoint detection
scaler = StandardScaler()
# Function to extract keypoints from a frame
def extract_keypoints(frame):
results = yolo_model(frame, verbose=False)
for r in results:
if r.keypoints is not None and len(r.keypoints) > 0:
keypoints = r.keypoints.xyn.tolist()[0] # Use the first person's keypoints
flattened_keypoints = [kp for keypoint in keypoints for kp in keypoint[:2]] # Flatten x, y values
return flattened_keypoints
return None # Return None if no keypoints are detected
# Function to process each frame
def process_frame(frame):
results = yolo_model(frame, verbose=False)
for box in results[0].boxes:
cls = int(box.cls[0]) # Class ID
confidence = float(box.conf[0])
if cls == 0 and confidence > 0.5: # Detect persons only
x1, y1, x2, y2 = map(int, box.xyxy[0]) # Bounding box coordinates
# Extract ROI for classification
roi = frame[y1:y2, x1:x2]
if roi.size > 0:
keypoints = extract_keypoints(roi)
if keypoints is not None and len(keypoints) > 0:
keypoints_scaled = scaler.fit_transform([keypoints])
keypoints_reshaped = keypoints_scaled.reshape((1, 1, len(keypoints)))
prediction = (lstm_model.predict(keypoints_reshaped) > 0.5).astype(int)[0][0]
color = (0, 0, 255) if prediction == 1 else (0, 255, 0)
label = 'Suspicious' if prediction == 1 else 'Normal'
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
return frame
# Gradio video streaming function
def video_processing(video_frame):
frame = cv2.cvtColor(video_frame, cv2.COLOR_BGR2RGB) # Convert to RGB
processed_frame = process_frame(frame)
return processed_frame
# Launch Gradio app
gr.Interface(
fn=video_processing,
inputs=gr.Video(streaming=True), # Correct the Video component
outputs="video",
live=True,
title="Suspicious Activity Detection"
).launch(debug=True)