Add video file using Git LFS

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,4 +1,6 @@
 import gradio as gr
+import cv2
+import os
 from core.chessboard_detector import ChessboardDetector
 
 detector = ChessboardDetector(
@@ -7,8 +9,6 @@ detector = ChessboardDetector(
     full_classifier_model_path="onnx/layout_recognition/v1.onnx"
 )
 
-
-
 # 数据集路径
 dict_cate_names = {
     '.': '.',
@@ -33,7 +33,95 @@ dict_cate_names = {
 dict_cate_names_reverse = {v: k for k, v in dict_cate_names.items()}
 
 
+### 构建 examples 
+
+def build_examples():
+    examples = []
+    # 读取 examples 目录下的所有图片
+    for file in os.listdir("examples"):
+        if file.endswith(".jpg"):
+            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])
+
+    return 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="""
         .image img {
             max-height: 512px;
@@ -43,12 +131,21 @@ with gr.Blocks(
     gr.Markdown("""
                 ## 棋盘检测, 棋子识别
 
+                x 表示 有遮挡位置   
+                . 表示 棋盘上的普通交叉点  
+
                 步骤:  
                     1. 流程分成两步，第一步检测边缘  
                     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")
@@ -83,18 +180,30 @@ with gr.Blocks(
                 visible=True,
             )
 
+    with gr.Row():
+        with gr.Column():
+            gr.Examples(full_examples, inputs=[image_input, video_input], label="示例视频、图片")
+
 
     def detect_chessboard(image):
-        original_image_with_keypoints, transformed_image, cells_labels_str, scores, time_info = detector.pred_detect_board_and_classifier(image)
+        if image is None:
+            return None, None, None, None
+        
+        try:
+            original_image_with_keypoints, transformed_image, cells_labels_str, scores, time_info = detector.pred_detect_board_and_classifier(image)
 
-        # 将 cells_labels 转换为 DataFrame
-        # cells_labels 通过  \n 分割
-        annotation_10_rows = [item for item in cells_labels_str.split("\n")]
-        # 将 annotation_10_rows 转换成为 10 行 9 列的二维数组
-        annotation_arr_10_9 = [list(item) for item in annotation_10_rows]
+            # 将 cells_labels 转换为 DataFrame
+            # cells_labels ��过  \n 分割
+            annotation_10_rows = [item for item in cells_labels_str.split("\n")]
+            # 将 annotation_10_rows 转换成为 10 行 9 列的二维数组
+            annotation_arr_10_9 = [list(item) for item in annotation_10_rows]
 
-        # 将 棋子类别 转换为 中文
-        annotation_arr_10_9 = [[dict_cate_names_reverse[item] for item in row] for row in annotation_arr_10_9]
+            # 将 棋子类别 转换为 中文
+            annotation_arr_10_9 = [[dict_cate_names_reverse[item] for item in row] for row in annotation_arr_10_9]
+
+        except Exception as e:
+            gr.Warning(f"检测失败 图片或者视频布局错误")
+            return None, None, None, None
 
 
         return original_image_with_keypoints, transformed_image, annotation_arr_10_9, time_info
@@ -103,5 +212,9 @@ 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()