import cv2 import mediapipe as mp import numpy as np import tensorflow as tf from tensorflow.keras.layers import LSTM import streamlit as st labels = np.array(['FALL', 'LYING', 'SIT', 'STAND', 'MOVE']) n_time_steps = 25 mpPose = mp.solutions.pose pose = mpPose.Pose() mpDraw = mp.solutions.drawing_utils def custom_lstm(*args, **kwargs): kwargs.pop('time_major', None) return LSTM(*args, **kwargs) model = tf.keras.models.load_model('bro.h5', custom_objects={'LSTM': custom_lstm}) def make_landmark_timestep(results): c_lm = [] for id, lm in enumerate(results.pose_landmarks.landmark): c_lm.append(lm.x) c_lm.append(lm.y) c_lm.append(lm.z) c_lm.append(lm.visibility) return c_lm def draw_landmark_on_image(mpDraw, results, img, label): mpDraw.draw_landmarks(img, results.pose_landmarks, mpPose.POSE_CONNECTIONS) for id, lm in enumerate(results.pose_landmarks.landmark): h, w, c = img.shape cx, cy = int(lm.x * w), int(lm.y * h) if label != "FALL": cv2.circle(img, (cx, cy), 5, (0, 255, 0), cv2.FILLED) else: cv2.circle(img, (cx, cy), 5, (0, 0, 255), cv2.FILLED) return img def draw_class_on_image(label, img): font = cv2.FONT_HERSHEY_SIMPLEX bottomLeftCornerOfText = (10, 30) fontScale = 1 fontColor = (0, 255, 0) thickness = 2 lineType = 2 cv2.putText(img, label, bottomLeftCornerOfText, font, fontScale, fontColor, thickness, lineType) return img def detect(model, lm_list): lm_list = np.array(lm_list) lm_list = np.expand_dims(lm_list, axis=0) results = model.predict(lm_list) if results[0][0] >= 0.5: label = labels[0] elif results[0][1] >= 0.5: label = labels[1] elif results[0][2] >= 0.5: label = labels[2] elif results[0][3] >= 0.5: label = labels[3] elif results[0][4] >= 0.5: label = labels[4] else: label = "NONE DETECTION" return label def main(): st.title("Pose Detection and Classification") run_type = st.sidebar.selectbox("Select input type", ("Camera", "Video File")) if run_type == "Camera": cap = cv2.VideoCapture(0) else: video_file = st.sidebar.file_uploader("Upload a video", type=["mp4", "mov", "avi"]) if video_file is not None: # Temporarily save the uploaded video to disk to pass to cv2.VideoCapture with open("temp_video.mp4", "wb") as f: f.write(video_file.read()) cap = cv2.VideoCapture("temp_video.mp4") else: st.write("Please upload a video file.") return stframe = st.empty() label = 'Starting...' lm_list = [] while cap.isOpened(): success, img = cap.read() if not success: break imgRGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) results = pose.process(imgRGB) if results.pose_landmarks: c_lm = make_landmark_timestep(results) img = draw_landmark_on_image(mpDraw, results, img, label) img = draw_class_on_image(label, img) lm_list.append(c_lm) if len(lm_list) == n_time_steps: label = detect(model, lm_list) lm_list = [] stframe.image(img, channels="BGR") if cv2.waitKey(1) == ord('q'): break cap.release() if __name__ == '__main__': main()