DCT-Net/run_vid.py

43 lines
1.3 KiB
Python

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)