Split text prediction and attribution analysis into separate buttons

app.py

@@ -61,7 +61,6 @@ feature_extractor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-2
 
 # 定义推理函数
 def predict_text(image, text):
-    image = Image.fromarray(image)
     image_features = feature_extractor(images=image, return_tensors="pt")
 
     inputs = tokenizer.encode_plus(
@@ -84,7 +83,6 @@ def predict_text(image, text):
 
 # 定义归因分析函数
 def generate_attribution(image, text):
-    image = Image.fromarray(image)
     image_features = feature_extractor(images=image, return_tensors="pt")
 
     inputs = tokenizer.encode_plus(
@@ -130,20 +128,21 @@ def generate_attribution(image, text):
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(label="Input Image", interactive=True)
-            question_input = gr.Textbox(label="Question", lines=2, interactive=True)
+            input_image = gr.Image(label="Input Image", interactive=True, height=400)
+            question_input = gr.Textbox(label="Question", lines=3, max_lines=3)
             clear_button = gr.Button("Clear")
         with gr.Column():
             predict_button = gr.Button("Answer")
             prediction_output = gr.Textbox(label="Answer", lines=2, interactive=False)
             attribution_button = gr.Button("Generate Attribution")
             with gr.Row():
-                attribution_image_1 = gr.Image(label="Attribution Image", interactive=False)
-                attribution_image_2 = gr.Image(label="Attribution with Contours", interactive=False)
+                attribution_image_1 = gr.Image(label="Attribution Image", interactive=False, height=400)
+                attribution_image_2 = gr.Image(label="Attribution with Contours", interactive=False, height=400)
 
     # 按钮事件绑定
     predict_button.click(predict_text, inputs=[input_image, question_input], outputs=prediction_output)
     attribution_button.click(generate_attribution, inputs=[input_image, question_input], outputs=[attribution_image_1, attribution_image_2])
     clear_button.click(lambda: (None, "", ""), outputs=[input_image, question_input, prediction_output])
 
+# 启动 Gradio 界面
 demo.launch()