feat: 性能更新

HISTORY.md

@@ -1,3 +1,8 @@
+### 2025-01-10
+
+1. 移除 视频
+2. 增加图片
+
 ### 2025-01-05
 
 1. 使用 swinv2 tiny 模型

app.py

@@ -1,12 +1,12 @@
 import gradio as gr
-import cv2
+# import cv2
 import os
 from core.chessboard_detector import ChessboardDetector
 
 detector = ChessboardDetector(
-    det_model_path="onnx/det/v1.onnx", 
+    det_model_path="onnx/det/v2.onnx", 
     pose_model_path="onnx/pose/4_v2.onnx",
-    full_classifier_model_path="onnx/layout_recognition/v3.onnx"
+    full_classifier_model_path="onnx/layout_recognition/v5.onnx"
 )
 
 # 数据集路径
@@ -39,13 +39,9 @@ def build_examples():
     examples = []
     # 读取 examples 目录下的所有图片
     for file in os.listdir("examples"):
-        if file.endswith(".jpg"):
+        if file.endswith(".jpg") or file.endswith(".png"):
             image_path = os.path.join("examples", file)
-            examples.append([image_path, None])
-        
-        elif file.endswith(".mp4"):
-            video_path = os.path.join("examples", file)
-            examples.append([None, video_path])
+            examples.append([image_path])
 
     return examples
 
@@ -53,76 +49,7 @@ def build_examples():
 full_examples = build_examples()
 
 
-def get_video_frame_with_processs(video_data, process: str = '00:00') -> cv2.UMat:
-    """
-    获取视频指定位置的帧
-    """
-
-    # 读取视频
-    cap = cv2.VideoCapture(video_data)
-    if not cap.isOpened():
-        gr.Warning("无法打开视频")
-        return None
-    
-    # 获取视频的帧率
-    fps = cap.get(cv2.CAP_PROP_FPS)
-
-    # process 是 00:00
-    process_time = process.split(":")
-    minutes = int(process_time[0])
-    seconds = float(process_time[1])
-
-    # 计算总秒数
-    target_seconds = minutes * 60 + seconds
-
-    # 计算当前帧
-    current_frame = int(target_seconds * fps)
-
-    # 设置到指定帧
-    cap.set(cv2.CAP_PROP_POS_FRAMES, current_frame)
-
-    # 读取当前帧
-    ret, frame = cap.read()
-    cap.release()
-
-    if not ret:
-        gr.Warning("无法读取视频帧")
-        return None
-    
-    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-
-    return frame_rgb
-
-
-
-
-with gr.Blocks(
-    js="""
-        async () => {
-            document.addEventListener('timeupdate', function(e) {
-                // 检查事件源是否是视频元素
-                if (e.target.matches('#video_player video')) {
-                    const video = e.target;
-                    const currentTime = video.currentTime;
-                    // 转换成 00:00 格式
-                    let minutes = Math.floor(currentTime / 60);
-                    let seconds = Math.floor(currentTime % 60);
-                    let formattedTime = `${minutes.toString().padStart(2,'0')}:${seconds.toString().padStart(2,'0')}`;
-                    
-                    // 更新输入框值
-                    let processInput = document.querySelector("#video_process textarea");
-                    if(processInput) {
-                        processInput.value = formattedTime;
-                        processInput.text = formattedTime;
-
-                        processInput.dispatchEvent(new Event("input"));
-                    }
-
-                }
-            }, true);  // 使用捕获阶段
-        }
-    """,
-    css="""
+with gr.Blocks(css="""
         .image img {
             max-height: 512px;
         }
@@ -139,13 +66,6 @@ with gr.Blocks(
                     2. 对整个棋盘画面进行棋子分类预测
                 """
     )
-
-    with gr.Row():
-        with gr.Column():
-            video_input = gr.Video(label="上传视频", interactive=True, elem_id="video_player", height=356)
-            video_process = gr.Textbox(label="当前时间", interactive=True, elem_id="video_process",  value="00:00")
-            extract_frame_btn = gr.Button("从视频提取当前帧")
-
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(label="上传棋盘图片", type="numpy", elem_classes="image")
@@ -182,7 +102,7 @@ with gr.Blocks(
 
     with gr.Row():
         with gr.Column():
-            gr.Examples(full_examples, inputs=[image_input, video_input], label="示例视频、图片")
+            gr.Examples(full_examples, inputs=[image_input], label="示例视频、图片")
 
 
     def detect_chessboard(image):
@@ -212,9 +132,5 @@ with gr.Blocks(
                        inputs=[image_input], 
                        outputs=[original_image_with_keypoints, transformed_image, layout_pred_info, use_time])
 
-    extract_frame_btn.click(fn=get_video_frame_with_processs, 
-                            inputs=[video_input, video_process], 
-                            outputs=[image_input])
-
 if __name__ == "__main__":
     demo.launch()

core/runonnx/rtmpose.py

@@ -350,7 +350,12 @@ class RTMPOSE_ONNX(BaseONNX):
         
         return original_keypoints
     
-    def draw_pred(self, img: cv2.UMat, keypoints: np.ndarray, scores: np.ndarray, is_rgb: bool = True) -> cv2.UMat:
+    def draw_pred(self, 
+                  img: cv2.UMat, 
+                  keypoints: np.ndarray, 
+                  scores: np.ndarray, 
+                  is_rgb: bool = True,
+                  score_threshold: float = 0.6) -> cv2.UMat:
         """
         Draw the keypoints results on the image.
         """
@@ -361,14 +366,18 @@ class RTMPOSE_ONNX(BaseONNX):
         colors = self.bone_colors
 
         for i, (point, score) in enumerate(zip(keypoints, scores)):
-            if score > 0.3:  # 设置置信度阈值
+           
                 x, y = map(int, point)
                 # 使用不同颜色标注不同的关键点
                 color = colors[i]
 
                 cv2.circle(img, (x, y), 5, (int(color[0]), int(color[1]), int(color[2])), -1)
                 # 添加关键点索引标注
-                cv2.putText(img, self.bone_names[i], (x+5, y+5), 
+                if score < score_threshold:  # 设置置信度阈值
+                    text = f"{self.bone_names[i]}: {score:.2f}"
+                else:
+                    text = f"{self.bone_names[i]}"
+                cv2.putText(img, text, (x+5, y+5), 
                             cv2.FONT_HERSHEY_SIMPLEX, 1.0, (int(color[0]), int(color[1]), int(color[2])), 1)
                 
                 # 绘制 关节连接线
@@ -383,7 +392,7 @@ class RTMPOSE_ONNX(BaseONNX):
             link_color = colors[start_index]
 
             # 绘制连线
-            if scores[start_index] > 0.3 and scores[end_index] > 0.3:
+            if scores[start_index] > score_threshold and scores[end_index] > score_threshold:
                 start_point = tuple(map(int, start_keypoint))
                 end_point = tuple(map(int, end_keypoint))
                 cv2.line(img, start_point, end_point,