diff --git a/applications/tools/first-order-demo.py b/applications/tools/first-order-demo.py index a2075324eb2e2b87f40d81598963b70f9af62cf4..149f721be6ab2211900990c64f535d2bb7988d5f 100644 --- a/applications/tools/first-order-demo.py +++ b/applications/tools/first-order-demo.py @@ -25,6 +25,7 @@ parser.add_argument("--weight_path", 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 @@ if __name__ == "__main__": 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 index cfc9c4ec6613646ed59f67b20c7d403cf486a41b..794300272ee88de3fa66375d377254a2b05554e3 100644 --- a/ppgan/apps/first_order_predictor.py +++ b/ppgan/apps/first_order_predictor.py @@ -45,7 +45,8 @@ class FirstOrderPredictor(BasePredictor): 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 @@ class FirstOrderPredictor(BasePredictor): 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 @@ class FirstOrderPredictor(BasePredictor): 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)