Add output filename to support concurrent work in FOM. (PaddlePaddle#204

) 增加文件名的参数filename，以便可以多并发执行 first_order_predictor.py
mayite · Mar 3, 2021 · c11d914 · c11d914
1 parent 183ffd1
commit c11d914
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/applications/tools/first-order-demo.py b/applications/tools/first-order-demo.py
@@ -25,6 +25,7 @@
 parser.add_argument("--source_image", type=str, help="path to source image")
 parser.add_argument("--driving_video", type=str, help="path to driving video")
 parser.add_argument("--output", default='output', help="path to output")
+parser.add_argument("--filename", default='result.mp4', help="filename to output")
 parser.add_argument("--relative",
                     dest="relative",
                     action="store_true",
@@ -65,6 +66,7 @@
         paddle.set_device('cpu')
 
     predictor = FirstOrderPredictor(output=args.output,
+                                    filename=args.filename,
                                     weight_path=args.weight_path,
                                     config=args.config,
                                     relative=args.relative,

diff --git a/ppgan/apps/first_order_predictor.py b/ppgan/apps/first_order_predictor.py
@@ -45,7 +45,8 @@ def __init__(self,
                  adapt_scale=False,
                  find_best_frame=False,
                  best_frame=None,
-                 ratio=1.0):
+                 ratio=1.0,
+                 filename='result.mp4'):
         if config is not None and isinstance(config, str):
             self.cfg = yaml.load(config, Loader=yaml.SafeLoader)
         elif isinstance(config, dict):
@@ -88,6 +89,7 @@ def __init__(self,
         if not os.path.exists(output):
             os.makedirs(output)
         self.output = output
+        self.filename = filename
         self.relative = relative
         self.adapt_scale = adapt_scale
         self.find_best_frame = find_best_frame
@@ -170,7 +172,7 @@ def run(self, source_image, driving_video):
                 frame = cv2.copyTo(patch, mask, frame)
 
             out_frame.append(frame)
-        imageio.mimsave(os.path.join(self.output, 'result.mp4'),
+        imageio.mimsave(os.path.join(self.output, self.filename),
                         [frame for frame in out_frame],
                         fps=fps)