Update app.py

app.py

@@ -112,10 +112,17 @@ def process_image(model, tokenizer, transform, device, check_type, image, text):
     return image, current_vis, current_bpe
 
 # 事件处理函数
-def update_index(change):
-    
-    aa=1
-    return 1
+# 上一项和下一项按钮
+def update_index(direction, current_vis, current_bpe, current_index):
+    # 计算新的索引
+    new_index = max(0, min(current_index + direction, len(current_vis) - 1))
+
+    # 更新可视化内容
+    return (
+        current_vis[new_index], 
+        format_bpe_display(current_bpe[new_index]), 
+        new_index  # 更新索引
+    )
 
 def format_bpe_display(bpe):
     # 使用HTML标签来设置字体大小、颜色，加粗，并居中
@@ -162,6 +169,10 @@ with gr.Blocks(title="BPE Visualization Demo") as demo:
             
             bpe_display = gr.Markdown("Current BPE: ")
 
+    current_vis_state = gr.State([])
+    current_bpe_state = gr.State([])
+    current_index_state = gr.State(0)
+    
     # 事件处理
     @spaces.GPU
     def on_run_clicked(model_type, image, text):
@@ -174,23 +185,36 @@ with gr.Blocks(title="BPE Visualization Demo") as demo:
         print("current_vis",len(current_vis))
         print("current_bpe",len(current_bpe))
         
-        return image, current_vis[0],f"<div style='text-align:center; font-size:20px;'><strong>Current BPE: <span style='color:red;'>{current_bpe[0]}</span></strong></div>", gr.update(visible=True), gr.update(visible=True)
+        # return image, current_vis[0],f"<div style='text-align:center; font-size:20px;'><strong>Current BPE: <span style='color:red;'>{current_bpe[0]}</span></strong></div>", gr.update(visible=True), gr.update(visible=True)
+        return (
+            image, 
+            current_vis[current_index], 
+            format_bpe_display(current_bpe[current_index]), 
+            gr.update(visible=True), 
+            gr.update(visible=True), 
+            current_vis,  # 存储整个列表
+            current_bpe,  # 存储整个列表
+            current_index  # 存储当前索引
+        )
 
     
+    
     run_btn.click(
         on_run_clicked,
         inputs=[model_type, image_input, text_input],
-        outputs=[orig_img, heatmap, bpe_display, prev_btn, next_btn],  # 让它们显示
+        outputs=[orig_img, heatmap, bpe_display, prev_btn, next_btn, current_vis_state, current_bpe_state, current_index_state]
     )
     
     prev_btn.click(
-        lambda: (*update_index(-1), current_index),
-        outputs=[heatmap, bpe_display]
+        update_index,
+        inputs=[gr.State(-1), current_vis_state, current_bpe_state, current_index_state],
+        outputs=[heatmap, bpe_display, current_index_state]
     )
     
     next_btn.click(
-        lambda: (*update_index(1), current_index),
-        outputs=[heatmap, bpe_display]
+        update_index,
+        inputs=[gr.State(1), current_vis_state, current_bpe_state, current_index_state],
+        outputs=[heatmap, bpe_display, current_index_state]
     )