import torch import torch.nn as nn import torch.nn.functional as F from pydantic import BaseModel from ultralytics import YOLO import os import cv2 import numpy as np class Model(nn.Module): def __init__(self, A, nnode, nfeature, nclass): super().__init__() self.fc1 = nn.Linear(nnode * nfeature, 512) self.fc2 = nn.Linear(512, nclass) def forward(self, x): x = x.view(-1, int(x.size(1) * x.size(2))) x = F.relu(self.fc1(x)) x = F.dropout(x, 0.7, training=self.training) return self.fc2(x) def extract_keypoint(get_keypoint, keypoint): # nose nose_x, nose_y = keypoint[get_keypoint.NOSE] # eye # left_eye_x, left_eye_y = keypoint[get_keypoint.LEFT_EYE] # right_eye_x, right_eye_y = keypoint[get_keypoint.RIGHT_EYE] # # ear # left_ear_x, left_ear_y = keypoint[get_keypoint.LEFT_EAR] # right_ear_x, right_ear_y = keypoint[get_keypoint.RIGHT_EAR] # shoulder left_shoulder_x, left_shoulder_y = keypoint[get_keypoint.LEFT_SHOULDER] right_shoulder_x, right_shoulder_y = keypoint[get_keypoint.RIGHT_SHOULDER] # elbow left_elbow_x, left_elbow_y = keypoint[get_keypoint.LEFT_ELBOW] right_elbow_x, right_elbow_y = keypoint[get_keypoint.RIGHT_ELBOW] # wrist left_wrist_x, left_wrist_y = keypoint[get_keypoint.LEFT_WRIST] right_wrist_x, right_wrist_y = keypoint[get_keypoint.RIGHT_WRIST] # hip left_hip_x, left_hip_y = keypoint[get_keypoint.LEFT_HIP] right_hip_x, right_hip_y = keypoint[get_keypoint.RIGHT_HIP] # knee left_knee_x, left_knee_y = keypoint[get_keypoint.LEFT_KNEE] right_knee_x, right_knee_y = keypoint[get_keypoint.RIGHT_KNEE] # ankle left_ankle_x, left_ankle_y = keypoint[get_keypoint.LEFT_ANKLE] right_ankle_x, right_ankle_y = keypoint[get_keypoint.RIGHT_ANKLE] return [ nose_x, nose_y , left_shoulder_x, left_shoulder_y , right_shoulder_x, right_shoulder_y, left_elbow_x, left_elbow_y, right_elbow_x, right_elbow_y, left_wrist_x, left_wrist_y, right_wrist_x, right_wrist_y, left_hip_x, left_hip_y, right_hip_x, right_hip_y, left_knee_x, left_knee_y, right_knee_x, right_knee_y, left_ankle_x, left_ankle_y, right_ankle_x, right_ankle_y ] class GetKeypoint(BaseModel): NOSE: int = 0 LEFT_EYE: int = 1 RIGHT_EYE: int = 2 LEFT_EAR: int = 3 RIGHT_EAR: int = 4 LEFT_SHOULDER: int = 5 RIGHT_SHOULDER: int = 6 LEFT_ELBOW: int = 7 RIGHT_ELBOW: int = 8 LEFT_WRIST: int = 9 RIGHT_WRIST: int = 10 LEFT_HIP: int = 11 RIGHT_HIP: int = 12 LEFT_KNEE: int = 13 RIGHT_KNEE: int = 14 LEFT_ANKLE: int = 15 RIGHT_ANKLE: int = 16 class Box(BaseModel): left : int top : int right : int bottom : int box_conf : float pose_classifer_conf : float label : str class YoloPoseClassfier(object): def __init__(self, device_id, yolo_pose_model_path, pose_classfier_model_path, confidence_threshold=0.5): self.confidence_threshold = confidence_threshold self.device = torch.device("cpu") if device_id == -1 else torch.device(f"cuda:{device_id}") # define get key point function self.get_keypoint = GetKeypoint() # load model self.pose_classfier_model = Model(None, 13, 2, 2) self.pose_classfier_model.load_state_dict(torch.load(pose_classfier_model_path)) self.pose_classfier_model.eval() self.pose_classfier_model.to(self.device) self.yolo_pose_model = YOLO(yolo_pose_model_path) self.yolo_pose_model.to(self.device) # 新增躯干角度校验函数 def calculate_torso_angle(keypoints): # 获取肩膀（索引5、6）和臀部（索引11、12）关键点 left_shoulder = keypoints[5] right_shoulder = keypoints[6] left_hip = keypoints[11] right_hip = keypoints[12] # 计算肩膀中点与臀部中点的向量 shoulder_center = ((left_shoulder[0] + right_shoulder[0]) / 2, (left_shoulder[1] + right_shoulder[1]) / 2) hip_center = ((left_hip[0] + right_hip[0]) / 2, (left_hip[1] + right_hip[1]) / 2) # 计算与垂直方向的夹角（跌倒时角度接近90度） dx = hip_center[0] - shoulder_center[0] dy = hip_center[1] - shoulder_center[1] angle = np.degrees(np.arctan2(abs(dx), abs(dy))) return angle def __call__(self, image : np.ndarray): pose_classfier_results = [] yolo_pose_results = self.yolo_pose_model(image, conf=self.confidence_threshold) for result in yolo_pose_results[0]: boxes = result.boxes.xyxy.cpu().numpy().tolist() confs = result.boxes.conf.cpu().numpy().tolist() all_keypoints = result.keypoints.data.cpu().numpy().tolist() for box, conf, keypoints in zip(boxes, confs, all_keypoints): x1, y1, x2, y2 = box x, y, w, h = x1, y1, x2 - x1, y2 - y1 n_keypoints = [[(kp[0] - x) / w - 0.5, (kp[1] - y) / h - 0.5] if kp[0] > 0 and kp[1] > 0 else kp[:2] for kp in keypoints] n_keypoints = extract_keypoint(self.get_keypoint, n_keypoints) if n_keypoints[-12:].count(0) >= 2 * 2: continue if n_keypoints.count(0) >= 4 * 2: continue if w < h: continue pose_data = torch.Tensor([n_keypoints]).to(self.device) pose_data = pose_data.reshape(1, 13, 2) with torch.no_grad(): p = self.pose_classfier_model(pose_data) prob = F.softmax(p) index = prob.argmax() if index == 0: score = float(prob[0][index].cpu().numpy()) pose_classfier_results.append( Box(left=x1, top=y1, right=x2, bottom=y2, box_conf=conf, pose_classifer_conf=score, label="falling")) return pose_classfier_results if __name__ == "__main__": yolo_pose_model_path = "yolo11l-pose.pt" pose_classfier_model_path = "fall_cls.pth" model = YoloPoseClassfier(-1, yolo_pose_model_path, pose_classfier_model_path, 0.5) test_images_dir = r"C:\Users\86187\Desktop\fall_test" save_image_dir = "." for image_name in os.listdir(test_images_dir): print(image_name) image_path = os.path.join(test_images_dir, image_name) image = cv2.imread(image_path) pose_classfier_results = model(image) for res in pose_classfier_results: cv2.rectangle(image, (int(res.left), int(res.top)), (int(res.right), int(res.bottom)), (0, 255, 0), 2) cv2.imwrite(os.path.join(save_image_dir, image_name), image)