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)
def main():
if opt.mode == 'add':
if util.is_img(opt.media_path):
core.addmosaic_img(opt)
elif util.is_video(opt.media_path):
core.addmosaic_video(opt)
else:
print('This type of file is not supported')
elif opt.mode == 'clean':
if util.is_img(opt.media_path):
core.cleanmosaic_img(opt)
elif util.is_video(opt.media_path):
if opt.netG == 'video':
core.cleanmosaic_video_fusion(opt)
else:
core.cleanmosaic_video_byframe(opt)
else:
print('This type of file is not supported')
util.clean_tempfiles(tmp_init=False)
img_result = img_origin.copy()
img_result = (img_origin * (1 - mask) + img_tmp * mask).astype('uint8')
return img_result
netG = loadmodel.pix2pix(os.path.join(opt.model_dir, opt.model_name),
opt.model_type_netG,
use_gpu=opt.use_gpu)
net_mosaic_pos = loadmodel.unet(os.path.join(opt.model_dir,
opt.mosaic_position_model_name),
use_gpu=opt.use_gpu)
filepaths = util.Traversal(opt.input_dir)
for path in filepaths:
if util.is_img(path):
print('Clean Mosaic:', path)
img_origin = cv2.imread(path)
x, y, size = get_mosaic_position(img_origin)
img_result = img_origin.copy()
if size != 0:
img_mosaic = img_origin[y - size:y + size, x - size:x + size]
img_fake = runmodel.run_pix2pix(img_mosaic,
netG,
use_gpu=opt.use_gpu)
img_result = replace_mosaic(img_origin, img_fake, x, y, size)
cv2.imwrite(os.path.join(opt.result_dir, os.path.basename(path)),
img_result)
elif util.is_video(path):
util.clean_tempfiles()
def main():
if opt.mode == 'add':
net = loadmodel.unet(opt)
path = opt.media_path
if util.is_img(path):
print('Add Mosaic:', path)
img = impro.imread(path)
mask = runmodel.get_ROI_position(img, net, opt)[0]
img = mosaic.addmosaic(img, mask, opt)
cv2.imwrite(os.path.join(opt.result_dir, os.path.basename(path)),
img)
elif util.is_video(path):
util.clean_tempfiles()
fps = ffmpeg.get_video_infos(path)[0]
ffmpeg.video2voice(path, './tmp/voice_tmp.mp3')
ffmpeg.video2image(
path, './tmp/video2image/output_%05d.' + opt.tempimage_type)
imagepaths = os.listdir('./tmp/video2image')
imagepaths.sort()
# get position
positions = []
for imagepath in imagepaths:
imagepath = os.path.join('./tmp/video2image', imagepath)
print('Find ROI location:', imagepath)
img = impro.imread(imagepath)
mask, x, y, area = runmodel.get_ROI_position(img, net, opt)
positions.append([x, y, area])
cv2.imwrite(
os.path.join('./tmp/ROI_mask',
os.path.basename(imagepath)), mask)
print('Optimize ROI locations...')
mask_index = filt.position_medfilt(np.array(positions), 7)
# add mosaic
print('Add mosaic to images...')
for i in range(len(imagepaths)):
mask_path = os.path.join('./tmp/ROI_mask',
imagepaths[mask_index[i]])
mask = impro.imread(mask_path)
img = impro.imread(
os.path.join('./tmp/video2image', imagepaths[i]))
img = mosaic.addmosaic(img, mask, opt)
cv2.imwrite(
os.path.join('./tmp/addmosaic_image',
os.path.basename(imagepaths[i])), img)
ffmpeg.image2video(
fps, './tmp/addmosaic_image/output_%05d.' + opt.tempimage_type,
'./tmp/voice_tmp.mp3',
os.path.join(
opt.result_dir,
os.path.splitext(os.path.basename(path))[0] + '_add.mp4'))
elif opt.mode == 'clean':
netG = loadmodel.pix2pix(opt)
net_mosaic_pos = loadmodel.unet_clean(opt)
path = opt.media_path
if util.is_img(path):
print('Clean Mosaic:', path)
img_origin = impro.imread(path)
x, y, size = runmodel.get_mosaic_position(img_origin,
net_mosaic_pos, opt)
img_result = img_origin.copy()
if size != 0:
img_mosaic = img_origin[y - size:y + size, x - size:x + size]
img_fake = runmodel.run_pix2pix(img_mosaic, netG, opt)
img_result = impro.replace_mosaic(img_origin, img_fake, x, y,
size, opt.no_feather)
cv2.imwrite(os.path.join(opt.result_dir, os.path.basename(path)),
img_result)
elif util.is_video(path):
util.clean_tempfiles()
fps = ffmpeg.get_video_infos(path)[0]
ffmpeg.video2voice(path, './tmp/voice_tmp.mp3')
ffmpeg.video2image(
path, './tmp/video2image/output_%05d.' + opt.tempimage_type)
positions = []
imagepaths = os.listdir('./tmp/video2image')
imagepaths.sort()
# get position
for imagepath in imagepaths:
imagepath = os.path.join('./tmp/video2image', imagepath)
img_origin = impro.imread(imagepath)
x, y, size = runmodel.get_mosaic_position(
img_origin, net_mosaic_pos, opt)
positions.append([x, y, size])
print('Find mosaic location:', imagepath)
print('Optimize mosaic locations...')
positions = np.array(positions)
for i in range(3):
positions[:, i] = filt.medfilt(positions[:, i],
opt.medfilt_num)
# clean mosaic
for i, imagepath in enumerate(imagepaths, 0):
imagepath = os.path.join('./tmp/video2image', imagepath)
x, y, size = positions[i][0], positions[i][1], positions[i][2]
img_origin = impro.imread(imagepath)
img_result = img_origin.copy()
if size != 0:
img_mosaic = img_origin[y - size:y + size,
x - size:x + size]
img_fake = runmodel.run_pix2pix(img_mosaic, netG, opt)
img_result = impro.replace_mosaic(img_origin, img_fake, x,
y, size, opt.no_feather)
cv2.imwrite(
os.path.join('./tmp/replace_mosaic',
os.path.basename(imagepath)), img_result)
print('Clean Mosaic:', imagepath)
ffmpeg.image2video(
fps, './tmp/replace_mosaic/output_%05d.' + opt.tempimage_type,
'./tmp/voice_tmp.mp3',
os.path.join(
opt.result_dir,
os.path.splitext(os.path.basename(path))[0] +
'_clean.mp4'))
util.clean_tempfiles(tmp_init=False)
play_cnt+= 1
frame_cnt += 1
def timer(opt,timerQueueime):
while True:
t = 1.0/opt.fps
time.sleep(t)
timerQueueime.put(True)
opt = Options().getparse()
system_type = 'Linux'
if 'Windows' in platform.platform():
system_type = 'Windows'
#-------------------------------Media Init-------------------------------
if util.is_img(opt.media):
img = cv2.imread(opt.media)
h_media,w_media = img.shape[:2]
elif util.is_video(opt.media):
fps,endtime,h_media,w_media = ffmpeg.get_video_infos(opt.media)
if opt.frame_num == 0:
opt.frame_num = int(endtime*fps-5)
if opt.ori_fps == 0:
opt.ori_fps = fps
util.makedirs('./tmp')
else:
print('Can not load this file!')
#-------------------------------Image Shape Init-------------------------------
if opt.screen==1:
limw = 80;limh = 24