stream last processed frame

app.py

@@ -48,6 +48,7 @@ def infer(video_input, in_threshold=0.5, num_people="Single person", render_mesh
         frame += 1
         img, mesh_paths, smplx_paths = inferer.infer(original_img, in_threshold, frame, multi_person, not(render_mesh))
         video_output.write(img)
+        yield img, None, None, None
     cap.release()
     video_output.release()
     cv2.destroyAllWindows()
@@ -62,7 +63,7 @@ def infer(video_input, in_threshold=0.5, num_people="Single person", render_mesh
     os.makedirs(save_path_smplx, exist_ok= True)
     os.system(f'zip -r {save_mesh_file} {save_path_mesh}')
     os.system(f'zip -r {save_smplx_file} {save_path_smplx}')
-    return video_path, save_mesh_file, save_smplx_file
+    yield img, video_path, save_mesh_file, save_smplx_file
 
 TITLE = '''<h1 align="center">SMPLer-X: Scaling Up Expressive Human Pose and Shape Estimation</h1>'''
 VIDEO = '''
@@ -108,12 +109,13 @@ with gr.Blocks(title="SMPLer-X", css=".gradio-container") as demo:
 
     with gr.Row():
         with gr.Column():
+            processed_frames = gr.Image(label="last frame")
             video_output = gr.Video(elem_classes="video")
         with gr.Column():
             meshes_output = gr.File(label="3D meshes")
             smplx_output = gr.File(label= "SMPL-X models")
     # example_images = gr.Examples([])
-    send_button.click(fn=infer, inputs=[video_input, threshold, num_people, mesh_as_vertices], outputs=[video_output, meshes_output, smplx_output])
+    send_button.click(fn=infer, inputs=[video_input, threshold, num_people, mesh_as_vertices], outputs=[processed_frames, video_output, meshes_output, smplx_output])
     # with gr.Row():
     example_videos = gr.Examples([
         ['/home/user/app/assets/01.mp4'],