update v3

app.py

@@ -1,11 +1,14 @@
 import gradio as gr
 # import cv2
 import os
+import base64
+from pathlib import Path
+    
 from core.chessboard_detector import ChessboardDetector
 
 detector = ChessboardDetector(
-    pose_model_path="onnx/pose/4_v3.onnx",
-    full_classifier_model_path="onnx/layout_recognition/nano_v1.onnx"
+    pose_model_path="onnx/pose/4_v6-0301.onnx",
+    full_classifier_model_path="onnx/layout_recognition/nano_v3-0319.onnx"
 )
 
 # 数据集路径
@@ -29,9 +32,51 @@ dict_cate_names = {
     '黑卒': 'p',
 }
 
+# 数据集路径
+dict_cate_images = {
+    'K': 'red_K.png',
+    'A': 'red_A.png',
+    'B': 'red_B.png',
+    'N': 'red_N.png',
+    'R': 'red_R.png',
+    'C': 'red_C.png',
+    'P': 'red_P.png',
+    'k': 'black_k.png',
+    'a': 'black_a.png',
+    'b': 'black_b.png',
+    'n': 'black_n.png',
+    'r': 'black_r.png',
+    'c': 'black_c.png',
+    'p': 'black_p.png',
+}
+
 dict_cate_names_reverse = {v: k for k, v in dict_cate_names.items()}
 
 
+ # 缓存图片的 base64 编码
+image_base64_cache = {}
+
+def get_image_base64(img_path):
+    if img_path in image_base64_cache:
+        return image_base64_cache[img_path]
+    
+    try:
+        img_full_path = Path("resources") / img_path
+        if not img_full_path.exists():
+            return ""
+            
+        with open(img_full_path, "rb") as img_file:
+            encoded = base64.b64encode(img_file.read()).decode('utf-8')
+            data_url = f"data:image/png;base64,{encoded}"
+            image_base64_cache[img_path] = data_url
+            return data_url
+    except Exception as e:
+        print(f"Error loading image {img_path}: {e}")
+        return 
+
+
+
+
 ### 构建 examples 
 
 def build_examples():
@@ -58,6 +103,7 @@ with gr.Blocks(css="""
                 ## 棋盘检测, 棋子识别
 
                 features: 轻量化模型
+                
 
                 x 表示 有遮挡位置   
                 . 表示 棋盘上的普通交叉点  
@@ -65,6 +111,9 @@ with gr.Blocks(css="""
                 步骤:  
                     1. 流程分成两步，第一步 keypoints 检测
                     2. 拉伸棋盘，并预测棋子
+                    
+                log:
+                    1. 优化棋子识别，增加对游戏棋盘的识别
                 """
     )
     with gr.Row():
@@ -95,16 +144,60 @@ with gr.Blocks(css="""
                 interactive=False,
                 visible=True,
             )
-            layout_pred_info = gr.Dataframe(
-                label="棋子识别",
-                interactive=False,
-                visible=True,
-            )
+            
+            # 添加 手风琴
+            with gr.Accordion("文字识别", open=False):
+                layout_pred_info = gr.Dataframe(
+                    label="棋子识别",
+                    interactive=False,
+                    visible=True,
+                )
+            
+            with gr.Accordion("棋子识别", open=True):
+                # 10 行 9 列的表格
+                table_html = gr.HTML(
+                    """
+                    <table>
+                    </table>
+                    """
+                )
 
     with gr.Row():
         with gr.Column():
             gr.Examples(
-                full_examples, inputs=[image_input], label="示例图片",  examples_per_page=15,)
+                full_examples[:10], inputs=[image_input], label="示例图片1",  examples_per_page=10,)
+            
+            gr.Examples(
+                full_examples[10:20], inputs=[image_input], label="示例图片2",  examples_per_page=10,)
+            
+            gr.Examples(
+                full_examples[20:], inputs=[image_input], label="示例图片3",  examples_per_page=10,)
+            
+
+    def gen_table_html(annotation_arr_10_9):
+        # 生成表格 HTML
+        html = "<table border='1' style='margin: auto;'>"
+        
+        for row in annotation_arr_10_9:
+            html += "<tr>"
+            for cell in row:
+                if cell == '.':
+                    # 普通交叉点
+                    html += "<td style='width: 60px; height: 60px; text-align: center;'></td>"
+                elif cell == 'x':
+                    # 遮挡位置
+                    html += "<td style='width: 60px; height: 60px; text-align: center;'>x</td>"
+                else:
+                    # 获取对应的图片文件名
+                    img_file = dict_cate_images.get(cell, '')
+                    img_data_base64 = get_image_base64(img_file)
+                    # 生成图片标签
+                    html += f"<td style='width: 60px; height: 60px; text-align: center; padding: 0;'><img src='{img_data_base64}' width='58' height='58'></td>"
+            html += "</tr>"
+        
+        html += "</table>"
+        return html
+        
 
 
     def detect_chessboard(image):
@@ -118,21 +211,23 @@ with gr.Blocks(css="""
             # 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_short = [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_short]
 
         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
+        table_html = gen_table_html(annotation_arr_10_9_short)
+
+        return original_image_with_keypoints, transformed_image, annotation_arr_10_9, table_html, time_info
 
     image_input.change(fn=detect_chessboard, 
                        inputs=[image_input], 
-                       outputs=[original_image_with_keypoints, transformed_image, layout_pred_info, use_time])
+                       outputs=[original_image_with_keypoints, transformed_image, layout_pred_info, table_html, use_time])
 
 if __name__ == "__main__":
     demo.launch()

onnx/layout_recognition/nano_v3-0319.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da66ba9809f15127f8ae729b1755e42ee61c100c4f9979ce0ef13602ac471298
+size 31101356