File size: 2,135 Bytes
713bd0b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
# model.py
import tensorflow as tf
from ultralytics import YOLO
import numpy as np
from sklearn.preprocessing import StandardScaler
class SuspiciousActivityModel:
def __init__(self, lstm_model_path, yolo_model_path):
# Load YOLO model
self.yolo_model = YOLO(yolo_model_path)
# Load LSTM model
self.lstm_model = tf.keras.models.load_model(lstm_model_path)
self.scaler = StandardScaler()
def extract_keypoints(self, frame):
"""
Extracts normalized keypoints from a frame using YOLO pose model.
"""
results = self.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]
flattened_keypoints = [kp for keypoint in keypoints for kp in keypoint[:2]]
return flattened_keypoints
return None
def process_frame(self, frame):
results = self.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:
x1, y1, x2, y2 = map(int, box.xyxy[0])
# Extract ROI for classification
roi = frame[y1:y2, x1:x2]
if roi.size > 0:
keypoints = self.extract_keypoints(roi)
if keypoints is not None and len(keypoints) > 0:
# Standardize and reshape keypoints for LSTM input
keypoints_scaled = self.scaler.fit_transform([keypoints])
keypoints_reshaped = keypoints_scaled.reshape((1, 1, len(keypoints)))
# Predict with LSTM model
prediction = (self.lstm_model.predict(keypoints_reshaped) > 0.5).astype(int)[0][0]
# Return label
return 'Suspicious' if prediction == 1 else 'Normal'
return 'Normal'
def detect_activity(self, frame):
return self.process_frame(frame)
|