def main():
if os.path.isdir(opt.media_path):
files = util.Traversal(opt.media_path)
else:
files = [opt.media_path]
if opt.mode == 'add':
netS = loadmodel.bisenet(opt, 'roi')
for file in files:
opt.media_path = file
if util.is_img(file):
core.addmosaic_img(opt, netS)
elif util.is_video(file):
core.addmosaic_video(opt, netS)
util.clean_tempfiles(opt, tmp_init=False)
else:
print('This type of file is not supported')
util.clean_tempfiles(opt, tmp_init=False)
elif opt.mode == 'clean':
netM = loadmodel.bisenet(opt, 'mosaic')
if opt.traditional:
netG = None
elif opt.netG == 'video':
netG = loadmodel.video(opt)
else:
netG = loadmodel.pix2pix(opt)
for file in files:
opt.media_path = file
if util.is_img(file):
core.cleanmosaic_img(opt, netG, netM)
elif util.is_video(file):
if opt.netG == 'video' and not opt.traditional:
core.cleanmosaic_video_fusion(opt, netG, netM)
else:
core.cleanmosaic_video_byframe(opt, netG, netM)
util.clean_tempfiles(opt, tmp_init=False)
else:
print('This type of file is not supported')
elif opt.mode == 'style':
netG = loadmodel.style(opt)
for file in files:
opt.media_path = file
if util.is_img(file):
core.styletransfer_img(opt, netG)
elif util.is_video(file):
core.styletransfer_video(opt, netG)
util.clean_tempfiles(opt, tmp_init=False)
else:
print('This type of file is not supported')
util.clean_tempfiles(opt, tmp_init=False)
if not os.path.isdir(opt.temp_dir):
util.file_init(opt)
def saveScriptModel(model,example,savepath):
model.cpu()
traced_script_module = torch.jit.trace(model, example)
# try ScriptModel
output = traced_script_module(example)
print(output)
traced_script_module.save(savepath)
savedir = '../cpp/res/models/'
util.makedirs(savedir)
opt.mosaic_position_model_path = '../pretrained_models/mosaic/mosaic_position.pth'
model = loadmodel.bisenet(opt,'mosaic')
example = torch.ones((1,3,360,360))
saveScriptModel(model,example,os.path.join(savedir,'mosaic_position.pt'))
# def main():
# if os.path.isdir(opt.media_path):
# files = util.Traversal(opt.media_path)
# else:
# files = [opt.media_path]
# if opt.mode == 'add':
# netS = loadmodel.bisenet(opt,'roi')
# for file in files:
# opt.media_path = file
default=96,
help='minimal roi size')
opt.parser.add_argument('--no_sclectscene', action='store_true', help='')
opt = opt.getparse()
util.makedirs(opt.savedir)
util.writelog(
os.path.join(opt.savedir, 'opt.txt'),
str(time.asctime(time.localtime(time.time()))) + '\n' + util.opt2str(opt))
videopaths = util.Traversal(opt.datadir)
videopaths = util.is_videos(videopaths)
random.shuffle(videopaths)
# def network
net = loadmodel.bisenet(opt, 'roi')
result_cnt = opt.startcnt
video_cnt = 1
starttime = datetime.datetime.now()
for videopath in videopaths:
try:
if opt.no_sclectscene:
timestamps = ['00:00:00']
else:
timestamps = []
fps, endtime, height, width = ffmpeg.get_video_infos(videopath)
for cut_point in range(1, int(
(endtime - opt.time) / opt.interval)):
util.clean_tempfiles(opt)
ffmpeg.video2image(