mirror of https://github.com/menyifang/DCT-Net
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
2 years ago
|
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)
|