From a8d79f781bd526210eadf30a0cb446adc064fd3b Mon Sep 17 00:00:00 2001 From: hypox64 Date: Tue, 26 May 2020 16:03:00 +0800 Subject: [PATCH] Temp --- cores/core.py | 76 +++++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 61 insertions(+), 15 deletions(-) diff --git a/cores/core.py b/cores/core.py index ee5b99d..6b2c0b1 100644 --- a/cores/core.py +++ b/cores/core.py @@ -37,18 +37,30 @@ def addmosaic_video(opt,netS): fps,imagepaths = video_init(opt,path)[:2] # get position positions = [] + t1 = time.time() + if not opt.no_preview: + cv2.namedWindow('preview', cv2.WINDOW_NORMAL) + + print('Find ROI location:') for i,imagepath in enumerate(imagepaths,1): img = impro.imread(os.path.join('./tmp/video2image',imagepath)) mask,x,y,size,area = runmodel.get_ROI_position(img,netS,opt) positions.append([x,y,area]) - cv2.imwrite(os.path.join('./tmp/ROI_mask',imagepath),mask) - print('Find ROI location:') - print('\r',str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='') + cv2.imwrite(os.path.join('./tmp/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(len(imagepaths)),util.get_bar(100*i/len(imagepaths),util.counttime(t1,t2,i,len(imagepaths)),num=35),end='') + print('\nOptimize ROI locations...') mask_index = filt.position_medfilt(np.array(positions), 7) # add mosaic print('Add Mosaic:') + t1 = time.time() 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])) @@ -58,8 +70,17 @@ def addmosaic_video(opt,netS): except Exception as e: print('Warning:',e) cv2.imwrite(os.path.join('./tmp/addmosaic_image',imagepaths[i]),img) - print('\r',str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='') + + #preview result and print + if not opt.no_preview: + cv2.imshow('preview',mask) + cv2.waitKey(1) & 0xFF + t2 = time.time() + print('\r',str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),util.counttime(t1,t2,i,len(imagepaths)),end='') + print() + if not opt.no_preview: + cv2.destroyAllWindows() ffmpeg.image2video( fps, './tmp/addmosaic_image/output_%06d.'+opt.tempimage_type, './tmp/voice_tmp.mp3', @@ -80,12 +101,26 @@ def styletransfer_video(opt,netG): positions = [] fps,imagepaths = video_init(opt,path)[:2] print('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('./tmp/video2image',imagepath)) img = runmodel.run_styletransfer(opt, netG, img) cv2.imwrite(os.path.join('./tmp/style_transfer',imagepath),img) - print('\r',str(i)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='') + + #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,len(imagepaths)),end='') + print() + if not opt.no_preview: + cv2.destroyAllWindows() suffix = os.path.basename(opt.model_path).replace('.pth','').replace('style_','') ffmpeg.image2video( fps, './tmp/style_transfer/output_%06d.'+opt.tempimage_type, @@ -101,24 +136,28 @@ def get_mosaic_positions(opt,netM,imagepaths,savemask=True): t1 = time.time() if not opt.no_preview: cv2.namedWindow('mosaic mask', cv2.WINDOW_NORMAL) + print('Find mosaic location:') 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) + positions.append([x,y,size]) + if savemask: + cv2.imwrite(os.path.join('./tmp/mosaic_mask',imagepath), mask) + + #preview result and print if not opt.no_preview: cv2.imshow('mosaic mask',mask) cv2.waitKey(1) & 0xFF - if savemask: - cv2.imwrite(os.path.join('./tmp/mosaic_mask',imagepath), mask) - positions.append([x,y,size]) 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) - np.save('./positions.npy', positions) + return positions def cleanmosaic_img(opt,netG,netM): @@ -150,6 +189,7 @@ def cleanmosaic_video_byframe(opt,netG,netM): # clean mosaic print('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('./tmp/video2image',imagepath)) @@ -166,16 +206,17 @@ def cleanmosaic_video_byframe(opt,netG,netM): except Exception as e: print('Warning:',e) cv2.imwrite(os.path.join('./tmp/replace_mosaic',imagepath),img_result) - #preview 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(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),util.counttime(t1,t2,i+1,len(imagepaths)),end='') + 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() - # to video ffmpeg.image2video( fps, './tmp/replace_mosaic/output_%06d.'+opt.tempimage_type, './tmp/voice_tmp.mp3', @@ -190,10 +231,14 @@ def cleanmosaic_video_fusion(opt,netG,netM): 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:') 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] @@ -211,18 +256,19 @@ def cleanmosaic_video_fusion(opt,netG,netM): if size>100: try:#Avoid unknown errors #reshape to network input shape - 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) + 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('./tmp/replace_mosaic',imagepath),img_result) + print('\r',str(i+1)+'/'+str(len(imagepaths)),util.get_bar(100*i/len(imagepaths),num=35),end='') print() ffmpeg.image2video( fps, -- GitLab