Update app.py

app.py

@@ -34,13 +34,14 @@ def image_vision(image_input_path, prompt):
     return_dict = model.predict_forward(**input_dict)
     print(return_dict)
     answer = return_dict["prediction"] # the text format answer
+    seg_image = return_dict["prediction_masks"]
     
-    return answer
+    return answer, seg_image
 
 def main_infer(image_input_path, prompt):
 
-    response = image_vision(image_input_path, prompt)
-    return response
+    answer, seg_image = image_vision(image_input_path, prompt)
+    return answer, seg_image
 
 # Gradio UI
 
@@ -55,11 +56,12 @@ with gr.Blocks() as demo:
                     submit_btn = gr.Button("Submit", scale=1)
             with gr.Column():
                 output_res = gr.Textbox(label="Response")
+                output_image = gr.Image(label="Segmentation")
 
     submit_btn.click(
         fn = main_infer,
         inputs = [image_input, instruction],
-        outputs = [output_res]
+        outputs = [output_res, output_image]
     )
 
 demo.queue().launch(show_api=False, show_error=True)