update

README.md

@@ -44,6 +44,7 @@ https://img.shields.io/badge/ModelScope-Spaces-blue)](https://modelscope.cn/#/mo
 * python 3
 * tensorflow (>=1.14)
 * easydict
+* imageio[ffmpeg]
 * numpy
 * both CPU/GPU are supported
 
@@ -101,7 +102,7 @@ python run.py
 video can be directly processed as image sequences, style choice [option: anime, 3d, handdrawn, sketch, artstyle, sd-design, sd-illustration]
 
 ```bash
-python run_vid.py --style anime
+python run_vid.py --video_path input.mp4 --save_path res.mp4 --style anime
 ```
 
 

run_vid.py

@@ -0,0 +1,42 @@
+import cv2, argparse
+import imageio
+from tqdm import tqdm
+import numpy as np
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
+
+def process(args):
+    style = args.style
+    print('choose style %s'%style)
+
+    reader = imageio.get_reader(args.video_path)
+    fps = reader.get_meta_data()['fps']
+    writer = imageio.get_writer(args.save_path, mode='I', fps=fps, codec='libx264')
+
+    if style == "anime":
+        style = ""
+    else:
+        style = '-' + style
+
+    model_name = 'damo/cv_unet_person-image-cartoon' + style + '_compound-models'
+    img_cartoon = pipeline(Tasks.image_portrait_stylization, model=model_name)
+
+    for _, img in tqdm(enumerate(reader)):
+        result = img_cartoon(img[..., ::-1])
+        res = result[OutputKeys.OUTPUT_IMG]
+        writer.append_data(res[..., ::-1].astype(np.uint8))
+    writer.close()
+    print('finished!')
+    print('result saved to %s'% args.save_path)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--video_path', type=str, default='PATH/TO/YOUR/MP4')
+    parser.add_argument('--save_path', type=str, default='res.mp4')
+    parser.add_argument('--style', type=str, default='anime')
+
+    args = parser.parse_args()
+
+    process(args)