def addmosaic_video(opt,netS): path = opt.media_path fps,imagepaths = video_init(opt,path) # get position positions = [] for i,imagepath in enumerate(imagepaths,1): img = impro.imread(os.path.join('./tmp/video2image',imagepath)) mask,x,y,area = runmodel.get_ROI_position(img,netS,opt) positions.append([x,y,area]) cv2.imwrite(os.path.join('./tmp/ROI_mask',imagepath),mask) print('\r','Find ROI location:'+str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=40),end='') print('\nOptimize ROI locations...') mask_index = filt.position_medfilt(np.array(positions), 7) # add mosaic for i in range(len(imagepaths)): mask = impro.imread(os.path.join('./tmp/ROI_mask',imagepaths[mask_index[i]]),'gray') img = impro.imread(os.path.join('./tmp/video2image',imagepaths[i])) if impro.mask_area(mask)>100: img = mosaic.addmosaic(img, mask, opt) cv2.imwrite(os.path.join('./tmp/addmosaic_image',imagepaths[i]),img) print('\r','Add Mosaic:'+str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=40),end='') print() 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'))
def cleanmosaic_video_byframe(opt,netG,netM): path = opt.media_path fps,imagepaths = video_init(opt,path) positions = [] # get position for i,imagepath in enumerate(imagepaths,1): img_origin = impro.imread(os.path.join('./tmp/video2image',imagepath)) x,y,size = runmodel.get_mosaic_position(img_origin,netM,opt)[:3] positions.append([x,y,size]) print('\r','Find mosaic location:'+str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=40),end='') print('\nOptimize 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): x,y,size = positions[i][0],positions[i][1],positions[i][2] img_origin = impro.imread(os.path.join('./tmp/video2image',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',imagepath),img_result) print('\r','Clean Mosaic:'+str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=40),end='') print() 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'))
def cleanmosaic_video_fusion(opt,netG,netM): path = opt.media_path N = 25 INPUT_SIZE = 128 fps,imagepaths = video_init(opt,path) positions = [] # get position for i,imagepath in enumerate(imagepaths,1): img_origin = impro.imread(os.path.join('./tmp/video2image',imagepath)) # x,y,size = runmodel.get_mosaic_position(img_origin,net_mosaic_pos,opt)[:3] x,y,size,mask = runmodel.get_mosaic_position(img_origin,netM,opt) cv2.imwrite(os.path.join('./tmp/mosaic_mask',imagepath), mask) positions.append([x,y,size]) print('\r','Find mosaic location:'+str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=40),end='') print('\nOptimize 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): x,y,size = positions[i][0],positions[i][1],positions[i][2] img_origin = impro.imread(os.path.join('./tmp/video2image',imagepath)) mask = cv2.imread(os.path.join('./tmp/mosaic_mask',imagepath),0) if size==0: cv2.imwrite(os.path.join('./tmp/replace_mosaic',imagepath),img_origin) else: mosaic_input = np.zeros((INPUT_SIZE,INPUT_SIZE,3*N+1), dtype='uint8') for j in range(0,N): img = impro.imread(os.path.join('./tmp/video2image',imagepaths[np.clip(i+j-12,0,len(imagepaths)-1)])) img = img[y-size:y+size,x-size:x+size] img = impro.resize(img,INPUT_SIZE) mosaic_input[:,:,j*3:(j+1)*3] = img mask = impro.resize(mask,np.min(img_origin.shape[:2])) mask = mask[y-size:y+size,x-size:x+size] mask = impro.resize(mask, INPUT_SIZE) mosaic_input[:,:,-1] = mask mosaic_input = data.im2tensor(mosaic_input,bgr2rgb=False,use_gpu=opt.use_gpu,use_transform = False,is0_1 = False) unmosaic_pred = netG(mosaic_input) #unmosaic_pred = (unmosaic_pred.cpu().detach().numpy()*255)[0] #img_fake = unmosaic_pred.transpose((1, 2, 0)) img_fake = data.tensor2im(unmosaic_pred,rgb2bgr = False ,is0_1 = False) img_result = impro.replace_mosaic(img_origin,img_fake,x,y,size,opt.no_feather) cv2.imwrite(os.path.join('./tmp/replace_mosaic',imagepath),img_result) print('\r','Clean Mosaic:'+str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=40),end='') print() 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'))
def get_mosaic_positions(opt, netM, imagepaths, savemask=True): # get mosaic position positions = [] t1 = time.time() if not opt.no_preview: cv2.namedWindow('mosaic mask', cv2.WINDOW_NORMAL) print('Step:2/4 -- Find mosaic location') for i, imagepath in enumerate(imagepaths, 1): img_origin = impro.imread( os.path.join(opt.temp_dir + '/video2image', imagepath)) x, y, size, mask = runmodel.get_mosaic_position(img_origin, netM, opt) positions.append([x, y, size]) if savemask: cv2.imwrite(os.path.join(opt.temp_dir + '/mosaic_mask', imagepath), mask) #preview result and print if not opt.no_preview: cv2.imshow('mosaic mask', mask) cv2.waitKey(1) & 0xFF t2 = time.time() print('\r', str(i) + '/' + str(len(imagepaths)), util.get_bar(100 * i / len(imagepaths), num=35), util.counttime(t1, t2, i, len(imagepaths)), end='') if not opt.no_preview: cv2.destroyAllWindows() print('\nOptimize mosaic locations...') positions = np.array(positions) for i in range(3): positions[:, i] = filt.medfilt(positions[:, i], opt.medfilt_num) return positions
def addmosaic_video(opt, netS): path = opt.media_path fps, imagepaths = video_init(opt, path)[:2] length = len(imagepaths) start_frame = int(imagepaths[0][7:13]) mask_index = get_roi_positions(opt, netS, imagepaths)[(start_frame - 1):] t1 = time.time() if not opt.no_preview: cv2.namedWindow('preview', cv2.WINDOW_NORMAL) # add mosaic print('Step:3/4 -- Add Mosaic:') t1 = time.time() # print(mask_index) for i, imagepath in enumerate(imagepaths, 1): mask = impro.imread( os.path.join( opt.temp_dir + '/ROI_mask', imagepaths[np.clip(mask_index[i - 1] - start_frame, 0, 1000000)]), 'gray') img = impro.imread( os.path.join(opt.temp_dir + '/video2image', imagepath)) if impro.mask_area(mask) > 100: try: #Avoid unknown errors img = mosaic.addmosaic(img, mask, opt) except Exception as e: print('Warning:', e) t = Thread(target=cv2.imwrite, args=(os.path.join(opt.temp_dir + '/addmosaic_image', imagepath), img)) t.start() os.remove(os.path.join(opt.temp_dir + '/video2image', imagepath)) #preview result and print if not opt.no_preview: cv2.imshow('preview', img) cv2.waitKey(1) & 0xFF t2 = time.time() print('\r', str(i) + '/' + str(length), util.get_bar(100 * i / length, num=35), util.counttime(t1, t2, i, length), end='') print() if not opt.no_preview: cv2.destroyAllWindows() print('Step:4/4 -- Convert images to video') ffmpeg.image2video( fps, opt.temp_dir + '/addmosaic_image/output_%06d.' + opt.tempimage_type, opt.temp_dir + '/voice_tmp.mp3', os.path.join(opt.result_dir, os.path.splitext(os.path.basename(path))[0] + '_add.mp4'))
def cleanmosaic_video_byframe(opt, netG, netM): path = opt.media_path fps, imagepaths = video_init(opt, path)[:2] positions = get_mosaic_positions(opt, netM, imagepaths, savemask=True) t1 = time.time() if not opt.no_preview: cv2.namedWindow('clean', cv2.WINDOW_NORMAL) # clean mosaic print('Step:3/4 -- Clean Mosaic:') length = len(imagepaths) for i, imagepath in enumerate(imagepaths, 0): x, y, size = positions[i][0], positions[i][1], positions[i][2] img_origin = impro.imread( os.path.join(opt.temp_dir + '/video2image', imagepath)) img_result = img_origin.copy() if size > 100: try: #Avoid unknown errors img_mosaic = img_origin[y - size:y + size, x - size:x + size] if opt.traditional: img_fake = runmodel.traditional_cleaner(img_mosaic, opt) else: img_fake = runmodel.run_pix2pix(img_mosaic, netG, opt) mask = cv2.imread( os.path.join(opt.temp_dir + '/mosaic_mask', imagepath), 0) img_result = impro.replace_mosaic(img_origin, img_fake, mask, x, y, size, opt.no_feather) except Exception as e: print('Warning:', e) cv2.imwrite(os.path.join(opt.temp_dir + '/replace_mosaic', imagepath), img_result) os.remove(os.path.join(opt.temp_dir + '/video2image', imagepath)) #preview result and print if not opt.no_preview: cv2.imshow('clean', img_result) cv2.waitKey(1) & 0xFF t2 = time.time() print('\r', str(i + 1) + '/' + str(length), util.get_bar(100 * i / length, num=35), util.counttime(t1, t2, i + 1, len(imagepaths)), end='') print() if not opt.no_preview: cv2.destroyAllWindows() print('Step:4/4 -- Convert images to video') ffmpeg.image2video( fps, opt.temp_dir + '/replace_mosaic/output_%06d.' + opt.tempimage_type, opt.temp_dir + '/voice_tmp.mp3', os.path.join( opt.result_dir, os.path.splitext(os.path.basename(path))[0] + '_clean.mp4'))
def get_mosaic_positions(opt,netM,imagepaths,savemask=True): # get mosaic position positions = [] for i,imagepath in enumerate(imagepaths,1): img_origin = impro.imread(os.path.join('./tmp/video2image',imagepath)) x,y,size,mask = runmodel.get_mosaic_position(img_origin,netM,opt) if savemask: cv2.imwrite(os.path.join('./tmp/mosaic_mask',imagepath), mask) positions.append([x,y,size]) print('\r','Find mosaic location:'+str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='') print('\nOptimize mosaic locations...') positions =np.array(positions) for i in range(3):positions[:,i] = filt.medfilt(positions[:,i],opt.medfilt_num) return positions
def styletransfer_video(opt,netG): path = opt.media_path positions = [] fps,imagepaths = video_init(opt,path) for i,imagepath in enumerate(imagepaths,1): img = impro.imread(os.path.join('./tmp/video2image',imagepath)) img = runmodel.run_styletransfer(opt, netG, img) cv2.imwrite(os.path.join('./tmp/style_transfer',imagepath),img) print('\r','Transfer:'+str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=40),end='') print() suffix = os.path.basename(opt.model_path).replace('.pth','').replace('style_','') ffmpeg.image2video( fps, './tmp/style_transfer/output_%05d.'+opt.tempimage_type, './tmp/voice_tmp.mp3', os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_'+suffix+'.mp4'))
def cleanmosaic_video_fusion(opt,netG,netM): path = opt.media_path N = 25 if 'HD' in os.path.basename(opt.model_path): INPUT_SIZE = 256 else: INPUT_SIZE = 128 fps,imagepaths,height,width = video_init(opt,path) positions = get_mosaic_positions(opt,netM,imagepaths,savemask=True) # clean mosaic img_pool = np.zeros((height,width,3*N), dtype='uint8') for i,imagepath in enumerate(imagepaths,0): x,y,size = positions[i][0],positions[i][1],positions[i][2] # image read stream mask = cv2.imread(os.path.join('./tmp/mosaic_mask',imagepath),0) if i==0 : for j in range(0,N): img_pool[:,:,j*3:(j+1)*3] = impro.imread(os.path.join('./tmp/video2image',imagepaths[np.clip(i+j-12,0,len(imagepaths)-1)])) else: img_pool[:,:,0:(N-1)*3] = img_pool[:,:,3:N*3] img_pool[:,:,(N-1)*3:] = impro.imread(os.path.join('./tmp/video2image',imagepaths[np.clip(i+12,0,len(imagepaths)-1)])) img_origin = img_pool[:,:,int((N-1)/2)*3:(int((N-1)/2)+1)*3] if size==0: # can not find mosaic, cv2.imwrite(os.path.join('./tmp/replace_mosaic',imagepath),img_origin) else: mosaic_input = np.zeros((INPUT_SIZE,INPUT_SIZE,3*N+1), dtype='uint8') mosaic_input[:,:,0:N*3] = impro.resize(img_pool[y-size:y+size,x-size:x+size,:], INPUT_SIZE) mask_input = impro.resize(mask,np.min(img_origin.shape[:2]))[y-size:y+size,x-size:x+size] mosaic_input[:,:,-1] = impro.resize(mask_input, INPUT_SIZE) mosaic_input = data.im2tensor(mosaic_input,bgr2rgb=False,use_gpu=opt.use_gpu,use_transform = False,is0_1 = False) unmosaic_pred = netG(mosaic_input) img_fake = data.tensor2im(unmosaic_pred,rgb2bgr = False ,is0_1 = False) img_result = impro.replace_mosaic(img_origin,img_fake,mask,x,y,size,opt.no_feather) cv2.imwrite(os.path.join('./tmp/replace_mosaic',imagepath),img_result) print('\r','Clean Mosaic:'+str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='') print() 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'))
def styletransfer_video(opt, netG): path = opt.media_path positions = [] fps, imagepaths = video_init(opt, path)[:2] print('Step:2/4 -- Transfer') t1 = time.time() if not opt.no_preview: cv2.namedWindow('preview', cv2.WINDOW_NORMAL) length = len(imagepaths) for i, imagepath in enumerate(imagepaths, 1): img = impro.imread( os.path.join(opt.temp_dir + '/video2image', imagepath)) img = runmodel.run_styletransfer(opt, netG, img) cv2.imwrite(os.path.join(opt.temp_dir + '/style_transfer', imagepath), img) os.remove(os.path.join(opt.temp_dir + '/video2image', imagepath)) #preview result and print if not opt.no_preview: cv2.imshow('preview', img) cv2.waitKey(1) & 0xFF t2 = time.time() print('\r', str(i) + '/' + str(length), util.get_bar(100 * i / length, num=35), util.counttime(t1, t2, i, len(imagepaths)), end='') print() if not opt.no_preview: cv2.destroyAllWindows() suffix = os.path.basename(opt.model_path).replace('.pth', '').replace( 'style_', '') print('Step:4/4 -- Convert images to video') ffmpeg.image2video( fps, opt.temp_dir + '/style_transfer/output_%06d.' + opt.tempimage_type, opt.temp_dir + '/voice_tmp.mp3', os.path.join( opt.result_dir, os.path.splitext(os.path.basename(path))[0] + '_' + suffix + '.mp4'))
def cleanmosaic_video_byframe(opt,netG,netM): path = opt.media_path fps,imagepaths = video_init(opt,path)[:2] positions = get_mosaic_positions(opt,netM,imagepaths,savemask=False) # clean mosaic for i,imagepath in enumerate(imagepaths,0): x,y,size = positions[i][0],positions[i][1],positions[i][2] img_origin = impro.imread(os.path.join('./tmp/video2image',imagepath)) img_result = img_origin.copy() if size != 0: img_mosaic = img_origin[y-size:y+size,x-size:x+size] if opt.traditional: img_fake = runmodel.traditional_cleaner(img_mosaic,opt) else: 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',imagepath),img_result) print('\r','Clean Mosaic:'+str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='') print() 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'))
def cleanmosaic_video_fusion(opt, netG, netM): path = opt.media_path N = 25 if 'HD' in os.path.basename(opt.model_path): INPUT_SIZE = 256 else: INPUT_SIZE = 128 fps, imagepaths, height, width = video_init(opt, path) positions = get_mosaic_positions(opt, netM, imagepaths, savemask=True) t1 = time.time() if not opt.no_preview: cv2.namedWindow('clean', cv2.WINDOW_NORMAL) # clean mosaic print('Clean Mosaic:') length = len(imagepaths) img_pool = np.zeros((height, width, 3 * N), dtype='uint8') mosaic_input = np.zeros((INPUT_SIZE, INPUT_SIZE, 3 * N + 1), dtype='uint8') for i, imagepath in enumerate(imagepaths, 0): x, y, size = positions[i][0], positions[i][1], positions[i][2] # image read stream mask = cv2.imread( os.path.join(opt.temp_dir + '/mosaic_mask', imagepath), 0) if i == 0: for j in range(0, N): img_pool[:, :, j * 3:(j + 1) * 3] = impro.imread( os.path.join( opt.temp_dir + '/video2image', imagepaths[np.clip(i + j - 12, 0, len(imagepaths) - 1)])) else: img_pool[:, :, 0:(N - 1) * 3] = img_pool[:, :, 3:N * 3] img_pool[:, :, (N - 1) * 3:] = impro.imread( os.path.join( opt.temp_dir + '/video2image', imagepaths[np.clip(i + 12, 0, len(imagepaths) - 1)])) img_origin = img_pool[:, :, int((N - 1) / 2) * 3:(int((N - 1) / 2) + 1) * 3] img_result = img_origin.copy() if size > 100: try: #Avoid unknown errors #reshape to network input shape mosaic_input[:, :, 0:N * 3] = impro.resize( img_pool[y - size:y + size, x - size:x + size, :], INPUT_SIZE) mask_input = impro.resize(mask, np.min( img_origin.shape[:2]))[y - size:y + size, x - size:x + size] mosaic_input[:, :, -1] = impro.resize(mask_input, INPUT_SIZE) mosaic_input_tensor = data.im2tensor(mosaic_input, bgr2rgb=False, use_gpu=opt.use_gpu, use_transform=False, is0_1=False) unmosaic_pred = netG(mosaic_input_tensor) img_fake = data.tensor2im(unmosaic_pred, rgb2bgr=False, is0_1=False) img_result = impro.replace_mosaic(img_origin, img_fake, mask, x, y, size, opt.no_feather) except Exception as e: print('Warning:', e) cv2.imwrite(os.path.join(opt.temp_dir + '/replace_mosaic', imagepath), img_result) #preview result and print if not opt.no_preview: cv2.imshow('clean', img_result) cv2.waitKey(1) & 0xFF t2 = time.time() print('\r', str(i + 1) + '/' + str(length), util.get_bar(100 * i / length, num=35), util.counttime(t1, t2, i + 1, len(imagepaths)), end='') print() if not opt.no_preview: cv2.destroyAllWindows() ffmpeg.image2video( fps, opt.temp_dir + '/replace_mosaic/output_%06d.' + opt.tempimage_type, opt.temp_dir + '/voice_tmp.mp3', os.path.join( opt.result_dir, os.path.splitext(os.path.basename(path))[0] + '_clean.mp4'))
merge_img) if opt.savemask: cv2.imwrite( os.path.join(mask_save_path, opt.name + '%06d' % savecnt + '.png'), mask) # print("Processing:",imgpaths[i]," ","Remain:",len(imgpaths)*opt.fold-filecnt) # cv2.namedWindow('image', cv2.WINDOW_NORMAL) # cv2.imshow('image',img_mosaic) # cv2.waitKey(0) # cv2.destroyAllWindows() except Exception as e: print(imgpaths[i], e) if filecnt % 10 == 0: endtime = datetime.datetime.now() # used_time = (endtime-starttime).seconds used_time = (endtime - starttime).seconds all_length = len(imgpaths) * opt.fold percent = round(100 * filecnt / all_length, 1) all_time = used_time / filecnt * all_length print('\r', '', str(filecnt) + '/' + str(all_length) + ' ', util.get_bar(percent, 30), '', util.second2stamp(used_time) + '/' + util.second2stamp(all_time), 'f:' + str(savecnt), end=" ")
origindir = os.path.join(opt.savedir, savecnt, 'origin_image') maskdir = os.path.join(opt.savedir, savecnt, 'mask') util.makedirs(origindir) util.makedirs(maskdir) util.clean_tempfiles(opt) ffmpeg.video2image(videopath, opt.temp_dir + '/video2image/%05d.' + opt.tempimage_type, start_time=timestamp, last_time=util.second2stamp(opt.time)) endtime = datetime.datetime.now() print( str(video_cnt) + '/' + str(len(videopaths)) + ' ', util.get_bar(100 * video_cnt / len(videopaths), 35), '', util.second2stamp( (endtime - starttime).seconds) + '/' + util.second2stamp( (endtime - starttime).seconds / video_cnt * len(videopaths))) imagepaths = util.Traversal(opt.temp_dir + '/video2image') imagepaths = sorted(imagepaths) imgs = [] masks = [] # mask_flag = False # for imagepath in imagepaths: # img = impro.imread(imagepath) # mask = runmodel.get_ROI_position(img,net,opt,keepsize=True)[0] # imgs.append(img) # masks.append(mask)
def get_roi_positions(opt, netS, imagepaths, savemask=True): # resume continue_flag = False if os.path.isfile(os.path.join(opt.temp_dir, 'step.json')): step = util.loadjson(os.path.join(opt.temp_dir, 'step.json')) resume_frame = int(step['frame']) if int(step['step']) > 2: mask_index = np.load(os.path.join(opt.temp_dir, 'mask_index.npy')) return mask_index if int(step['step']) >= 2 and resume_frame > 0: pre_positions = np.load( os.path.join(opt.temp_dir, 'roi_positions.npy')) continue_flag = True imagepaths = imagepaths[resume_frame:] positions = [] t1 = time.time() if not opt.no_preview: cv2.namedWindow('mask', cv2.WINDOW_NORMAL) print('Step:2/4 -- Find mosaic location') img_read_pool = Queue(4) def loader(imagepaths): for imagepath in imagepaths: img_origin = impro.imread( os.path.join(opt.temp_dir + '/video2image', imagepath)) img_read_pool.put(img_origin) t = Thread(target=loader, args=(imagepaths, )) t.daemon = True t.start() for i, imagepath in enumerate(imagepaths, 1): img_origin = img_read_pool.get() mask, x, y, size, area = runmodel.get_ROI_position( img_origin, netS, opt) positions.append([x, y, area]) if savemask: t = Thread(target=cv2.imwrite, args=( os.path.join(opt.temp_dir + '/ROI_mask', imagepath), mask, )) t.start() if i % 1000 == 0: save_positions = np.array(positions) if continue_flag: save_positions = np.concatenate( (pre_positions, save_positions), axis=0) np.save(os.path.join(opt.temp_dir, 'roi_positions.npy'), save_positions) step = {'step': 2, 'frame': i + resume_frame} util.savejson(os.path.join(opt.temp_dir, 'step.json'), step) #preview result and print if not opt.no_preview: cv2.imshow('mask', mask) cv2.waitKey(1) & 0xFF t2 = time.time() print('\r', str(i) + '/' + str(len(imagepaths)), util.get_bar(100 * i / len(imagepaths), num=35), util.counttime(t1, t2, i, len(imagepaths)), end='') if not opt.no_preview: cv2.destroyAllWindows() print('\nOptimize ROI locations...') if continue_flag: positions = np.concatenate((pre_positions, positions), axis=0) mask_index = filt.position_medfilt(np.array(positions), 7) step = {'step': 3, 'frame': 0} util.savejson(os.path.join(opt.temp_dir, 'step.json'), step) np.save(os.path.join(opt.temp_dir, 'roi_positions.npy'), positions) np.save(os.path.join(opt.temp_dir, 'mask_index.npy'), np.array(mask_index)) return mask_index
def addmosaic_video(opt, netS): path = opt.media_path fps, imagepaths = video_init(opt, path)[:2] length = len(imagepaths) # get position positions = [] t1 = time.time() if not opt.no_preview: cv2.namedWindow('preview', cv2.WINDOW_NORMAL) print('Step:2/4 -- Find ROI location') for i, imagepath in enumerate(imagepaths, 1): img = impro.imread( os.path.join(opt.temp_dir + '/video2image', imagepath)) mask, x, y, size, area = runmodel.get_ROI_position(img, netS, opt) positions.append([x, y, area]) cv2.imwrite(os.path.join(opt.temp_dir + '/ROI_mask', imagepath), mask) #preview result and print if not opt.no_preview: cv2.imshow('preview', mask) cv2.waitKey(1) & 0xFF t2 = time.time() print('\r', str(i) + '/' + str(length), util.get_bar(100 * i / length, num=35), util.counttime(t1, t2, i, length), end='') print('\nOptimize ROI locations...') mask_index = filt.position_medfilt(np.array(positions), 7) # add mosaic print('Step:3/4 -- Add Mosaic:') t1 = time.time() for i, imagepath in enumerate(imagepaths, 1): mask = impro.imread( os.path.join(opt.temp_dir + '/ROI_mask', imagepaths[mask_index[i - 1]]), 'gray') img = impro.imread( os.path.join(opt.temp_dir + '/video2image', imagepath)) if impro.mask_area(mask) > 100: try: #Avoid unknown errors img = mosaic.addmosaic(img, mask, opt) except Exception as e: print('Warning:', e) cv2.imwrite(os.path.join(opt.temp_dir + '/addmosaic_image', imagepath), img) os.remove(os.path.join(opt.temp_dir + '/video2image', imagepath)) #preview result and print if not opt.no_preview: cv2.imshow('preview', img) cv2.waitKey(1) & 0xFF t2 = time.time() print('\r', str(i) + '/' + str(length), util.get_bar(100 * i / length, num=35), util.counttime(t1, t2, i, length), end='') print() if not opt.no_preview: cv2.destroyAllWindows() print('Step:4/4 -- Convert images to video') ffmpeg.image2video( fps, opt.temp_dir + '/addmosaic_image/output_%06d.' + opt.tempimage_type, opt.temp_dir + '/voice_tmp.mp3', os.path.join(opt.result_dir, os.path.splitext(os.path.basename(path))[0] + '_add.mp4'))
def cleanmosaic_video_fusion(opt, netG, netM): path = opt.media_path N, T, S = 2, 5, 3 LEFT_FRAME = (N * S) POOL_NUM = LEFT_FRAME * 2 + 1 INPUT_SIZE = 256 FRAME_POS = np.linspace(0, (T - 1) * S, T, dtype=np.int64) img_pool = [] previous_frame = None init_flag = True fps, imagepaths, height, width = video_init(opt, path) positions = get_mosaic_positions(opt, netM, imagepaths, savemask=True) t1 = time.time() if not opt.no_preview: cv2.namedWindow('clean', cv2.WINDOW_NORMAL) # clean mosaic print('Step:3/4 -- Clean Mosaic:') length = len(imagepaths) for i, imagepath in enumerate(imagepaths, 0): x, y, size = positions[i][0], positions[i][1], positions[i][2] input_stream = [] # image read stream if i == 0: # init for j in range(POOL_NUM): img_pool.append( impro.imread( os.path.join( opt.temp_dir + '/video2image', imagepaths[np.clip(i + j - LEFT_FRAME, 0, len(imagepaths) - 1)]))) else: # load next frame img_pool.pop(0) img_pool.append( impro.imread( os.path.join( opt.temp_dir + '/video2image', imagepaths[np.clip(i + LEFT_FRAME, 0, len(imagepaths) - 1)]))) img_origin = img_pool[LEFT_FRAME] img_result = img_origin.copy() if size > 50: try: #Avoid unknown errors for pos in FRAME_POS: input_stream.append( impro.resize( img_pool[pos][y - size:y + size, x - size:x + size], INPUT_SIZE)[:, :, ::-1]) if init_flag: init_flag = False previous_frame = input_stream[N] previous_frame = data.im2tensor(previous_frame, bgr2rgb=True, gpu_id=opt.gpu_id) input_stream = np.array(input_stream).reshape( 1, T, INPUT_SIZE, INPUT_SIZE, 3).transpose((0, 4, 1, 2, 3)) input_stream = data.to_tensor(data.normalize(input_stream), gpu_id=opt.gpu_id) with torch.no_grad(): unmosaic_pred = netG(input_stream, previous_frame) img_fake = data.tensor2im(unmosaic_pred, rgb2bgr=True) previous_frame = unmosaic_pred # previous_frame = data.tensor2im(unmosaic_pred,rgb2bgr = True) mask = cv2.imread( os.path.join(opt.temp_dir + '/mosaic_mask', imagepath), 0) img_result = impro.replace_mosaic(img_origin, img_fake, mask, x, y, size, opt.no_feather) except Exception as e: init_flag = True print('Error:', e) else: init_flag = True cv2.imwrite(os.path.join(opt.temp_dir + '/replace_mosaic', imagepath), img_result) os.remove(os.path.join(opt.temp_dir + '/video2image', imagepath)) #preview result and print if not opt.no_preview: cv2.imshow('clean', img_result) cv2.waitKey(1) & 0xFF t2 = time.time() print('\r', str(i + 1) + '/' + str(length), util.get_bar(100 * i / length, num=35), util.counttime(t1, t2, i + 1, len(imagepaths)), end='') print() if not opt.no_preview: cv2.destroyAllWindows() print('Step:4/4 -- Convert images to video') ffmpeg.image2video( fps, opt.temp_dir + '/replace_mosaic/output_%06d.' + opt.tempimage_type, opt.temp_dir + '/voice_tmp.mp3', os.path.join( opt.result_dir, os.path.splitext(os.path.basename(path))[0] + '_clean.mp4'))
if MOD == 'HD': #[128:384,128:384,:] --->256 cv2.imwrite( os.path.join(train_A_path, '%05d' % i + suffix + '.jpg'), mosaic_img) cv2.imwrite( os.path.join(train_B_path, '%05d' % i + suffix + '.jpg'), img) if MASK: cv2.imwrite( os.path.join(mask_path, '%05d' % i + suffix + '.png'), mask) elif MOD == 'pix2pix': merge_img = impro.makedataset(mosaic_img, img) cv2.imwrite(os.path.join(train_path, '%05d' % i + suffix + '.jpg'), merge_img) elif MOD == 'mosaic': cv2.imwrite( os.path.join(mosaic_path, '%05d' % i + suffix + '.jpg'), mosaic_img) cv2.imwrite(os.path.join(ori_path, '%05d' % i + suffix + '.jpg'), img) cv2.imwrite(os.path.join(mask_path, '%05d' % i + suffix + '.png'), mask) print('\r', 'Proc/all:' + str(i) + '/' + str(len(img_paths)), util.get_bar(100 * i / len(img_paths), num=40), end='') except Exception as e: print(img_path, e)
merge_img) if opt.savemask: cv2.imwrite( os.path.join(mask_save_path, opt.name + '%06d' % savecnt + '.png'), mask) # print("Processing:",imgpaths[i]," ","Remain:",len(imgpaths)*opt.fold-filecnt) # cv2.namedWindow('image', cv2.WINDOW_NORMAL) # cv2.imshow('image',img_mosaic) # cv2.waitKey(0) # cv2.destroyAllWindows() except Exception as e: print(imgpaths[i], e) if filecnt % 10 == 0: endtime = datetime.datetime.now() # used_time = (endtime-starttime).seconds used_time = (endtime - starttime).seconds all_length = len(imgpaths) * opt.fold percent = round(100 * filecnt / all_length, 1) all_time = used_time / filecnt * all_length print('\r', '', str(filecnt) + '/' + str(all_length) + ' ', util.get_bar(percent, 25), '', util.second2stamp(used_time) + '/' + util.second2stamp(all_time), 'f:' + str(savecnt), end=" ")
mod='gray') img_ori = impro.imread( os.path.join('./tmp/video2image', imagepaths[i])) img_mosaic = mosaic.addmosaic_normal(img_ori, mask, mosaic_size, model=models[mosaic_type], rect_rat=rect_rat) mask = impro.resize(mask, min(img_ori.shape[:2])) img_ori_crop = impro.resize( img_ori[y - size:y + size, x - size:x + size], 256) img_mosaic_crop = impro.resize( img_mosaic[y - size:y + size, x - size:x + size], 256) mask_crop = impro.resize( mask[y - size:y + size, x - size:x + size], 256) cv2.imwrite( os.path.join(ori_path, os.path.basename(imagepaths[i])), img_ori_crop) cv2.imwrite( os.path.join(mosaic_path, os.path.basename(imagepaths[i])), img_mosaic_crop) cv2.imwrite( os.path.join(mask_path, os.path.basename(imagepaths[i])), mask_crop) except Exception as e: print(e) print(util.get_bar(100 * i / len(videos), num=50))