From 2856640d40d4a378dcf2f2e4d4c4e3fdd0ad0e38 Mon Sep 17 00:00:00 2001 From: shuike Date: Tue, 9 Dec 2025 10:15:16 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E4=BA=BA=E8=84=B8=E8=A1=A8?= =?UTF-8?q?=E6=83=85=E5=8A=9F=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 3 + README.md | 25 +- ai_demo.py | 70 +++++ .../25c859cf-e15d-458d-82e9-81039c49f9ed.png | Bin .../7406b499-a8cb-4b9f-ad15-45558dd1602b.png | Bin util.py | 281 ++++++++++++++++++ weights/readme.md | 6 + 7 files changed, 383 insertions(+), 2 deletions(-) create mode 100755 ai_demo.py rename 25c859cf-e15d-458d-82e9-81039c49f9ed.png => doc/25c859cf-e15d-458d-82e9-81039c49f9ed.png (100%) rename 7406b499-a8cb-4b9f-ad15-45558dd1602b.png => doc/7406b499-a8cb-4b9f-ad15-45558dd1602b.png (100%) create mode 100755 util.py create mode 100644 weights/readme.md diff --git a/.gitignore b/.gitignore index 0dbf2f2..0638463 100644 --- a/.gitignore +++ b/.gitignore @@ -168,3 +168,6 @@ cython_debug/ # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ + +*.onnx +.DS_Store diff --git a/README.md b/README.md index c3de133..05c4f47 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,24 @@ -# student_server -学生学习用的测试服务器 \ No newline at end of file +## 测试环境 +conda activate py3.9 +pip install onnxruntime==1.14.0 +pip install opencv-python==4.5.5.64 + +## 目录结构如下 +. +├── __pycache__ +│ └── util.cpython-36.pyc +├── ai_demo.py +├── all_server_open（云）.py +├── db_con.py +├── doc +│ ├── 25c859cf-e15d-458d-82e9-81039c49f9ed.png +│ └── 7406b499-a8cb-4b9f-ad15-45558dd1602b.png +├── env.md +├── example_info_add.py +├── user_add.py +├── util.py +└── weights + ├── detection.onnx + ├── emotion.onnx + └── readme.md diff --git a/ai_demo.py b/ai_demo.py new file mode 100755 index 0000000..5f5519e --- /dev/null +++ b/ai_demo.py @@ -0,0 +1,70 @@ +from PIL import Image +from util import * +import json + +detector = FaceDetector('./weights/detection.onnx') +fer = HSEmotionRecognizer('./weights/emotion.onnx') + +def detect_face(frame): + boxes = detector.detect(frame, (640, 640)) + return boxes if boxes is not None and len(boxes) else None + + +def porcess_video(input_video_path): + stream = cv2.VideoCapture(input_video_path) + frame_width = int(stream.get(3)) + frame_height = int(stream.get(4)) + fps = stream.get(cv2.CAP_PROP_FPS) + + idx_class = fer.idx_to_class + frame_counter = 0 + emotion_reg_list = [] + emotion_score_list = [] + + while True: + ret, frame = stream.read() + if ret: + frame_counter += 1 + ## 每30帧进行处理一次 + if frame_counter % 30 == 0: + image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + boxes = detect_face(image) + if boxes is not None: + for box in boxes.astype('int32'): + x1, y1, x2, y2 = box[:4] + face_image = image[y1:y2, x1:x2] + pil_image = Image.fromarray(face_image).convert('RGB') + pil_image = pil_image.resize((224, 224)) # Adjust size as per model requirements + emotion, scores,pred = fer.predict_emotions(np.array(pil_image), logits=False) + ###### 做应用逻辑 + print("frame_id:",frame_counter) + print("emotion",emotion) + print("scores",scores[pred]) + emotion_reg_list.append(emotion) + emotion_score_list.append(float(scores[pred])) + print("idx_class",idx_class) + ###### + draw_emotion_bars(fer, frame, scores, (x2 + 10, y1), bar_height=15, width=100) + else: + print("退出") + break + + stream.release() + + + data = { + "video_url":input_video_path, + "emotion_reg_list":emotion_reg_list, + "emotion_score_list":emotion_score_list + } + json_str = json.dumps(data) + print("json_str:",json_str) + return json_str + +if __name__ == '__main__': + #自己录制一段视频 + video_path = "/Users/apple/Downloads/教学代码/111111.mp4" + porcess_video(video_path) + +#结果如下 +# json_str: {"video_url": "/Users/apple/Downloads/\u6559\u5b66\u4ee3\u7801/111111.mp4", "emotion_reg_list": ["NEUTRAL", "HAPPINESS"], "emotion_score_list": [0.5329781174659729, 0.43945103883743286]} \ No newline at end of file diff --git a/25c859cf-e15d-458d-82e9-81039c49f9ed.png b/doc/25c859cf-e15d-458d-82e9-81039c49f9ed.png similarity index 100% rename from 25c859cf-e15d-458d-82e9-81039c49f9ed.png rename to doc/25c859cf-e15d-458d-82e9-81039c49f9ed.png diff --git a/7406b499-a8cb-4b9f-ad15-45558dd1602b.png b/doc/7406b499-a8cb-4b9f-ad15-45558dd1602b.png similarity index 100% rename from 7406b499-a8cb-4b9f-ad15-45558dd1602b.png rename to doc/7406b499-a8cb-4b9f-ad15-45558dd1602b.png diff --git a/util.py b/util.py new file mode 100755 index 0000000..24ef503 --- /dev/null +++ b/util.py @@ -0,0 +1,281 @@ +import os +import cv2 +import numpy as np +import onnxruntime as ort + + +def distance2box(points, distance, max_shape=None): + x1 = points[:, 0] - distance[:, 0] + y1 = points[:, 1] - distance[:, 1] + x2 = points[:, 0] + distance[:, 2] + y2 = points[:, 1] + distance[:, 3] + if max_shape is not None: + x1 = x1.clamp(min=0, max=max_shape[1]) + y1 = y1.clamp(min=0, max=max_shape[0]) + x2 = x2.clamp(min=0, max=max_shape[1]) + y2 = y2.clamp(min=0, max=max_shape[0]) + return np.stack([x1, y1, x2, y2], axis=-1) + + +def distance2kps(points, distance, max_shape=None): + outputs = [] + for i in range(0, distance.shape[1], 2): + p_x = points[:, i % 2] + distance[:, i] + p_y = points[:, i % 2 + 1] + distance[:, i + 1] + if max_shape is not None: + p_x = p_x.clamp(min=0, max=max_shape[1]) + p_y = p_y.clamp(min=0, max=max_shape[0]) + outputs.append(p_x) + outputs.append(p_y) + return np.stack(outputs, axis=-1) + + +class FaceDetector: + def __init__(self, onnx_path=None, session=None): + from onnxruntime import InferenceSession + self.session = session + + self.batched = False + if self.session is None: + assert onnx_path is not None + assert os.path.exists(onnx_path) + self.session = InferenceSession(onnx_path, + providers=['CUDAExecutionProvider']) + self.nms_thresh = 0.4 + self.center_cache = {} + input_cfg = self.session.get_inputs()[0] + input_shape = input_cfg.shape + if isinstance(input_shape[2], str): + self.input_size = None + else: + self.input_size = tuple(input_shape[2:4][::-1]) + input_name = input_cfg.name + outputs = self.session.get_outputs() + if len(outputs[0].shape) == 3: + self.batched = True + output_names = [] + for output in outputs: + output_names.append(output.name) + self.input_name = input_name + self.output_names = output_names + self.use_kps = False + self._num_anchors = 1 + if len(outputs) == 6: + self.fmc = 3 + self._feat_stride_fpn = [8, 16, 32] + self._num_anchors = 2 + elif len(outputs) == 9: + self.fmc = 3 + self._feat_stride_fpn = [8, 16, 32] + self._num_anchors = 2 + self.use_kps = True + elif len(outputs) == 10: + self.fmc = 5 + self._feat_stride_fpn = [8, 16, 32, 64, 128] + self._num_anchors = 1 + elif len(outputs) == 15: + self.fmc = 5 + self._feat_stride_fpn = [8, 16, 32, 64, 128] + self._num_anchors = 1 + self.use_kps = True + + def forward(self, x, score_thresh): + scores_list = [] + bboxes_list = [] + points_list = [] + input_size = tuple(x.shape[0:2][::-1]) + blob = cv2.dnn.blobFromImage(x, + 1.0 / 128, + input_size, + (127.5, 127.5, 127.5), swapRB=True) + outputs = self.session.run(self.output_names, {self.input_name: blob}) + input_height = blob.shape[2] + input_width = blob.shape[3] + fmc = self.fmc + for idx, stride in enumerate(self._feat_stride_fpn): + if self.batched: + scores = outputs[idx][0] + boxes = outputs[idx + fmc][0] + boxes = boxes * stride + else: + scores = outputs[idx] + boxes = outputs[idx + fmc] + boxes = boxes * stride + + height = input_height // stride + width = input_width // stride + key = (height, width, stride) + if key in self.center_cache: + anchor_centers = self.center_cache[key] + else: + anchor_centers = np.stack(np.mgrid[:height, :width][::-1], axis=-1) + anchor_centers = anchor_centers.astype(np.float32) + + anchor_centers = (anchor_centers * stride).reshape((-1, 2)) + if self._num_anchors > 1: + anchor_centers = np.stack([anchor_centers] * self._num_anchors, axis=1) + anchor_centers = anchor_centers.reshape((-1, 2)) + if len(self.center_cache) < 100: + self.center_cache[key] = anchor_centers + + pos_indices = np.where(scores >= score_thresh)[0] + bboxes = distance2box(anchor_centers, boxes) + pos_scores = scores[pos_indices] + pos_bboxes = bboxes[pos_indices] + scores_list.append(pos_scores) + bboxes_list.append(pos_bboxes) + return scores_list, bboxes_list + + def detect(self, image, input_size=None, score_threshold=0.5, max_num=0, metric='default'): + assert input_size is not None or self.input_size is not None + input_size = self.input_size if input_size is None else input_size + image_ratio = float(image.shape[0]) / image.shape[1] + model_ratio = float(input_size[1]) / input_size[0] + if image_ratio > model_ratio: + new_height = input_size[1] + new_width = int(new_height / image_ratio) + else: + new_width = input_size[0] + new_height = int(new_width * image_ratio) + det_scale = float(new_height) / image.shape[0] + resized_img = cv2.resize(image, (new_width, new_height)) + det_img = np.zeros((input_size[1], input_size[0], 3), dtype=np.uint8) + det_img[:new_height, :new_width, :] = resized_img + + scores_list, bboxes_list = self.forward(det_img, score_threshold) + + scores = np.vstack(scores_list) + scores_ravel = scores.ravel() + order = scores_ravel.argsort()[::-1] + bboxes = np.vstack(bboxes_list) / det_scale + pre_det = np.hstack((bboxes, scores)).astype(np.float32, copy=False) + pre_det = pre_det[order, :] + keep = self.nms(pre_det) + det = pre_det[keep, :] + if 0 < max_num < det.shape[0]: + area = (det[:, 2] - det[:, 0]) * (det[:, 3] - det[:, 1]) + img_center = image.shape[0] // 2, image.shape[1] // 2 + offsets = np.vstack([(det[:, 0] + det[:, 2]) / 2 - img_center[1], + (det[:, 1] + det[:, 3]) / 2 - img_center[0]]) + offset_dist_squared = np.sum(np.power(offsets, 2.0), 0) + if metric == 'max': + values = area + else: + values = area - offset_dist_squared * 2.0 # some extra weight on the centering + index = np.argsort(values)[::-1] # some extra weight on the centering + index = index[0:max_num] + det = det[index, :] + return det + + def nms(self, outputs): + thresh = self.nms_thresh + x1 = outputs[:, 0] + y1 = outputs[:, 1] + x2 = outputs[:, 2] + y2 = outputs[:, 3] + scores = outputs[:, 4] + + order = scores.argsort()[::-1] + areas = (x2 - x1 + 1) * (y2 - y1 + 1) + + keep = [] + while order.size > 0: + i = order[0] + keep.append(i) + xx1 = np.maximum(x1[i], x1[order[1:]]) + yy1 = np.maximum(y1[i], y1[order[1:]]) + xx2 = np.minimum(x2[i], x2[order[1:]]) + yy2 = np.minimum(y2[i], y2[order[1:]]) + + w = np.maximum(0.0, xx2 - xx1 + 1) + h = np.maximum(0.0, yy2 - yy1 + 1) + inter = w * h + ovr = inter / (areas[i] + areas[order[1:]] - inter) + + indices = np.where(ovr <= thresh)[0] + order = order[indices + 1] + + return keep + + +class HSEmotionRecognizer: + def __init__(self, path): + self.img_size = 260 + self.mean = [0.485, 0.456, 0.406] + self.std = [0.229, 0.224, 0.225] + + self.idx_to_class = {0: 'ANGER', + 1: 'DISGUST', + 2: 'FEAR', + 3: 'HAPPINESS', + 4: 'NEUTRAL', + 5: 'SADNESS', + 6: 'SURPRISE'} + + self.ort_session = ort.InferenceSession(path, providers=['CUDAExecutionProvider']) + + def preprocess(self, img): + x = cv2.resize(img, (self.img_size, self.img_size)) / 255 + for i in range(3): + x[..., i] = (x[..., i] - self.mean[i]) / self.std[i] + return x.transpose(2, 0, 1).astype("float32")[np.newaxis, ...] + + def predict_emotions(self, face_img, logits=True): + scores = self.ort_session.run(None, {"input": self.preprocess(face_img)})[0][0] + x = scores + pred = np.argmax(x) + if not logits: + e_x = np.exp(x - np.max(x)[np.newaxis]) + e_x = e_x / e_x.sum()[None] + scores = e_x + return self.idx_to_class[pred], scores,pred + + def predict_multi_emotions(self, face_img_list, logits=True): + images = np.concatenate([self.preprocess(face_img) for face_img in face_img_list], axis=0) + scores = self.ort_session.run(None, {"input": images})[0] + pred = np.argmax(scores, axis=1) + x = scores + + if not logits: + e_x = np.exp(x - np.max(x, axis=1)[:, np.newaxis]) + e_x = e_x / e_x.sum(axis=1)[:, None] + scores = e_x + + return [self.idx_to_class[pred] for pred in pred], scores + + +def draw_rounded_bar(image, top_left, bottom_right, color, radius): + top_left = (int(top_left[0]), int(top_left[1])) + bottom_right = (int(bottom_right[0]), int(bottom_right[1])) + image_height, image_width = image.shape[:2] + + bottom_right = (min(bottom_right[0], image_width), min(bottom_right[1], image_height)) + + cv2.rectangle(image, (top_left[0] + radius, top_left[1]), (bottom_right[0] - radius, bottom_right[1]), color, + thickness=2) + cv2.rectangle(image, (top_left[0], top_left[1] + radius), (bottom_right[0], bottom_right[1] - radius), color, + thickness=2) + + cv2.circle(image, (top_left[0] + radius, top_left[1] + radius), radius, color, -1) + cv2.circle(image, (bottom_right[0] - radius, top_left[1] + radius), radius, color, -1) + cv2.circle(image, (top_left[0] + radius, bottom_right[1] - radius), radius, color, -1) + cv2.circle(image, (bottom_right[0] - radius, bottom_right[1] - radius), radius, color, -1) + + +def draw_emotion_bars(model, frame, scores, top_left, bar_height=25, width=150, spacing=10): + emotions = list(model.idx_to_class.values()) + start_x, start_y = top_left + + overlay = frame.copy() + overlay_height = len(emotions) * (bar_height + spacing) - spacing + cv2.rectangle(overlay, (start_x, start_y), (start_x + width, start_y + overlay_height), (0, 0, 0, 128), -1) + cv2.addWeighted(overlay, 0.4, frame, 1 - 0.4, 0, frame) + + for i, (emotion, score) in enumerate(zip(emotions, scores)): + bar_length = int(score * width) + bar_top_left = (start_x, start_y + i * (bar_height + spacing)) + bar_bottom_right = (start_x + bar_length, start_y + (i * (bar_height + spacing)) + bar_height) + draw_rounded_bar(frame, bar_top_left, bar_bottom_right, (0, 255, 0), radius=int(bar_height / 3)) + + text_position = (start_x + width + 5, bar_top_left[1] + bar_height // 2 + spacing // 4) + cv2.putText(frame, emotion, text_position, cv2.FONT_HERSHEY_COMPLEX, 0.4, (66, 15, 7), 1) diff --git a/weights/readme.md b/weights/readme.md new file mode 100644 index 0000000..34f216d --- /dev/null +++ b/weights/readme.md @@ -0,0 +1,6 @@ +请将onnx下载后,放入到weights文件夹下 + +http://m51c.dongtaiyuming.net:28620/directlink/5/aibox/models/detection.onnx + +http://m51c.dongtaiyuming.net:28620/directlink/5/aibox/models/emotion.onnx +