def addmosaic_img(opt): net = loadmodel.unet(opt) path = opt.media_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.splitext(os.path.basename(path))[0] + '_add.jpg'), img)
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.unet(opt) 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) else: print('This type of file is not supported') elif opt.mode == 'clean': netM = loadmodel.unet_clean(opt) 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) 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) else: print('This type of file is not supported') util.clean_tempfiles(tmp_init=False)
def addmosaic_video(opt): net = loadmodel.unet(opt) path = opt.media_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: print('Find ROI location:', imagepath) img = impro.imread(os.path.join('./tmp/video2image', 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', 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 = impro.imread( os.path.join('./tmp/ROI_mask', imagepaths[mask_index[i]])) img = impro.imread(os.path.join('./tmp/video2image', imagepaths[i])) img = mosaic.addmosaic(img, mask, opt) cv2.imwrite(os.path.join('./tmp/addmosaic_image', 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'))
(255, 255, 255), -1) mask = (cv2.blur(mask, (eclosion_num, eclosion_num))) mask = mask / 255.0 img_tmp = np.zeros(img_origin.shape) img_tmp[y - size:y + size, x - size:x + size] = img_fake 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)
from models import runmodel,loadmodel from util import mosaic,util,ffmpeg from util import image_processing as impro from options.addmosaic_options import AddOptions opt = AddOptions().getparse() #find mosaic position in image and add mosaic to this image def add_mosaic_to_image(path): img = cv2.imread(path) mask =runmodel.run_unet_rectim(img,net,use_gpu = opt.use_gpu) mask = impro.mask_threshold(mask,opt.mask_extend,opt.mask_threshold) img = mosaic.addmosaic(img,mask,opt.mosaic_size,opt.output_size,model = opt.mosaic_mod) return img net = loadmodel.unet(os.path.join(opt.model_dir,opt.model_name),use_gpu = opt.use_gpu) filepaths = util.Traversal(opt.input_dir) for path in filepaths: if util.is_img(path): img = add_mosaic_to_image(path) 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) for imagepath in os.listdir('./tmp/video2image'): imagepath = os.path.join('./tmp/video2image',imagepath) print(imagepath)
import sys sys.path.append("..") from models import runmodel, loadmodel from util import mosaic, util, ffmpeg, filt from util import image_processing as impro from options import Options opt = Options().getparse() util.file_init(opt) videos = os.listdir('./video') videos.sort() opt.model_path = '../pretrained_models/add_youknow_128.pth' opt.use_gpu = True net = loadmodel.unet(opt) for path in videos: path = os.path.join('./video', path) util.clean_tempfiles() 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 = [] img_ori_example = impro.imread( os.path.join('./tmp/video2image', imagepaths[0])) mask_avg = np.zeros((impro.resize(img_ori_example, 128)).shape[:2])
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)