diff --git a/README.md b/README.md index 1fbf1e1..37c4684 100644 --- a/README.md +++ b/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 ``` diff --git a/input.mp4 b/input.mp4 new file mode 100644 index 0000000..f8f3428 Binary files /dev/null and b/input.mp4 differ diff --git a/run_vid.py b/run_vid.py new file mode 100644 index 0000000..4dd051e --- /dev/null +++ b/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)