Spaces:
Runtime error
Runtime error
| # 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) | |