filename

main.py

@@ -35,6 +35,25 @@ inferencers = {"Estimate human 2d poses":human, "Estimate human 2d hand poses":h
 
 print("[INFO]: Downloaded models!")
 
+def check_extension(video):
+    split_tup = os.path.splitext(video)
+
+    # extract the file name and extension
+    file_name = split_tup[0]
+    file_extension = split_tup[1]
+
+    if file_extension is not ".mp4":
+        clip = moviepy.VideoFileClip(video)
+
+        video = file_name+".mp4"
+        clip.write_videofile(video)
+    
+    return video
+
+
+
+
+
 def tracking(video, model, boxes=True):
     print("[INFO] Loading model...")
     # Load an official or custom model
@@ -82,10 +101,7 @@ def show_tracking(video_content):
 
 
 def pose3d(video):
-    clip = moviepy.VideoFileClip(video)
-
-    video = video[:-5]+".mp4"
-    clip.write_videofile(video)
+    video = check_extension(video)
 
 
     # Define new unique folder
@@ -107,12 +123,8 @@ def pose3d(video):
     return "".join(out_file)
 
 
-def pose2d(video):
-
-    clip = moviepy.VideoFileClip(video)
-
-    video = video[:-5]+".mp4"
-    clip.write_videofile(video)
+def pose2d(video, kpt_threshold):
+    video = check_extension(video)
 
     # Define new unique folder
     add_dir = str(uuid.uuid4())
@@ -124,7 +136,7 @@ def pose2d(video):
                             return_vis=True,
                             thickness=2,
                             rebase_keypoint_height=True,
-                            #kpt_thr=kpt_thr,
+                            kpt_thr=kpt_threshold,
                             device="cuda"
                             )    
     
@@ -135,12 +147,8 @@ def pose2d(video):
     return "".join(out_file)
 
 
-def pose2dhand(video):
-    clip = moviepy.VideoFileClip(video)
-
-    video = video[:-5]+".mp4"
-    clip.write_videofile(video)
-
+def pose2dhand(video, kpt_threshold):
+    video = check_extension(video)
 
     # Define new unique folder
     add_dir = str(uuid.uuid4())
@@ -152,6 +160,7 @@ def pose2dhand(video):
                                  return_vis=True,
                                  thickness=2,
                                  rebase_keypoint_height=True,
+                                 kpt_thr=kpt_threshold,
                                  device="cuda")    
     
     result = [result for result in result_generator] #next(result_generator)        
@@ -174,11 +183,11 @@ def run_UI():
                         submit_pose3d_file = gr.Button("Make 3d pose estimation")
                         submit_hand_file = gr.Button("Make 2d hand estimation")
                         submit_detect_file = gr.Button("Detect and track objects")
-
-                    video_output1 = gr.PlayableVideo(height=512,  label = "Estimate human 2d poses", show_label=True)
-                    video_output2 = gr.PlayableVideo(height=512,  label = "Estimate human 3d poses", show_label=True)
-                    video_output3 = gr.PlayableVideo(height=512,  label = "Estimate human hand poses", show_label=True)
-                    video_output4 = gr.Video(height=512)
+                    with gr.Column():
+                        video_output1 = gr.PlayableVideo(height=512,  label = "Estimate human 2d poses", show_label=True)
+                        video_output2 = gr.PlayableVideo(height=512,  label = "Estimate human 3d poses", show_label=True)
+                        video_output3 = gr.PlayableVideo(height=512,  label = "Estimate human hand poses", show_label=True)
+                        video_output4 = gr.Video(height=512)
 
             with gr.Tab("Record video with webcam"):
                 with gr.Row():
@@ -201,7 +210,7 @@ def run_UI():
         
         # From file
         submit_pose_file.click(fn=pose2d, 
-                            inputs= video_input, 
+                            inputs=  [video_input, file_kpthr], 
                             outputs = video_output1)
         
         submit_pose3d_file.click(fn=pose3d, 
@@ -209,7 +218,7 @@ def run_UI():
                                 outputs = video_output2)
         
         submit_hand_file.click(fn=pose2dhand, 
-                            inputs= video_input, 
+                            inputs= [video_input, file_kpthr], 
                             outputs = video_output3)
         
         submit_detect_file.click(fn=show_tracking, 
@@ -218,7 +227,7 @@ def run_UI():
         
         # Web
         submit_pose_web.click(fn=pose2d, 
-                            inputs = webcam_input, 
+                            inputs = [webcam_input, web_kpthr], 
                             outputs = webcam_output1)
         
         submit_pose3d_web.click(fn=pose3d, 
@@ -226,7 +235,7 @@ def run_UI():
                                 outputs = webcam_output2)
         
         submit_hand_web.click(fn=pose2dhand, 
-                            inputs= webcam_input, 
+                            inputs= [webcam_input, web_kpthr], 
                             outputs = webcam_output3)
         
         submit_detect_web.click(fn=show_tracking,