Update stf_utils.py

stf_utils.py

@@ -11,6 +11,8 @@ import imageio
 import numpy as np
 from rich.progress import track
 from tqdm import tqdm
+import zipfile
+import shutil
 
 import stf_alternative
 
@@ -75,20 +77,32 @@ class STFPipeline:
                  config_path: str = "front_config.json",
                  checkpoint_path: str = "089.pth",
                  #root_path: str = "works"
-                 root_path: str = "/tmp/works"
+                 root_path: str = "/tmp/works",
+                 female_video: bool=True
                  
     ):
         #os.makedirs(root_path, exist_ok=True)
-        import shutil; shutil.copytree('/home/user/app/stf/works', '/tmp/works', dirs_exist_ok=True)
+        shutil.copytree('/home/user/app/stf/works', '/tmp/works', dirs_exist_ok=True)
         
-        import zipfile
-        dir_zip='/tmp/works/preprocess/nasilhong_f_v1_front/crop_video_front_one_piece_dress_nodded_cut.zip'
-        dir_target='/tmp/works/preprocess/nasilhong_f_v1_front/'
-        zipfile.ZipFile(dir_zip, 'r').extractall(dir_target)
         
-        dir_zip='/tmp/works/preprocess/nasilhong_f_v1_front/front_one_piece_dress_nodded_cut.zip'
-        dir_target='/tmp/works/preprocess/nasilhong_f_v1_front/'
-        import zipfile; zipfile.ZipFile(dir_zip, 'r').extractall(dir_target)
+
+        if female_video:
+            dir_zip= os.path.join(root_path, 'preprocess/nasilhong_f_v1_front/crop_video_front_one_piece_dress_nodded_cut.zip')
+            dir_target=os.path.join(root_path,'preprocess/nasilhong_f_v1_front/')
+            zipfile.ZipFile(dir_zip, 'r').extractall(dir_target)
+            
+            dir_zip=os.path.join(root_path,'preprocess/nasilhong_f_v1_front/front_one_piece_dress_nodded_cut.zip')
+            dir_target=os.path.join(root_path,'preprocess/nasilhong_f_v1_front/')
+            zipfile.ZipFile(dir_zip, 'r').extractall(dir_target)
+        else:
+            dir_zip= os.path.join(root_path, 'preprocess/Ian_v3_front/crop_video_Cam2_2309071202_0012_Natural_Looped.zip')
+            dir_target=os.path.join(root_path,'preprocess/Ian_v3_front/')
+            zipfile.ZipFile(dir_zip, 'r').extractall(dir_target)
+            
+            dir_zip=os.path.join(root_path,'preprocess/Ian_v3_front/Cam2_2309071202_0012_Natural_Looped.zip')
+            dir_target=os.path.join(root_path,'preprocess/Ian_v3_front/')
+            zipfile.ZipFile(dir_zip, 'r').extractall(dir_target)
+            
         
 
         self.config_path = os.path.join(stf_path, config_path)
@@ -173,7 +187,6 @@ class STFPipeline:
                 )
                 for idx, (it, chunk) in enumerate(gen_infer, pivot):
                     frame = next(reader)
-                    print('frame=', frame.shape, frame[:10])
                     composed = self.template.compose(idx, frame, it)
                     frame_name = f"{idx}".zfill(5)+".jpg"
                     results.append(it['pred'])
@@ -183,5 +196,8 @@ class STFPipeline:
 
         print('STFPipeline execute 3')
         images2video(results, save_path)
+
+        save_path_aud = save_path.replace('.mp4', '_aud.mp4')
+        merge_audio_video(save_path, audio, save_path_aud)
                                 
-        return save_path
+        return save_path_aud #save_path