From 6c6d052213038969261f5b116446354016f9c43d Mon Sep 17 00:00:00 2001 From: hypox64 Date: Fri, 16 Aug 2019 16:11:39 +0800 Subject: [PATCH] fix auto mosaic size --- README.md | 20 +++--- README_CN.md | 20 +++--- deepmosaic.py | 191 +++++++++++++++++++++++++------------------------ util/mosaic.py | 8 +-- 4 files changed, 122 insertions(+), 117 deletions(-) diff --git a/README.md b/README.md index 67f6c07..6347149 100755 --- a/README.md +++ b/README.md @@ -17,21 +17,21 @@ Download this version via [[Google Drive]](https://drive.google.com/open?id=1LTE ![image](./imgs/GUI.png)
Attentions:
-- Require Windows_x86_64, Windows10 is better.
-- Different pre-trained models are suitable for different effects.
-- Run time depends on computer performance.
-- If output video cannot be played, you can try with [potplayer](https://daumpotplayer.com/download/). + - Require Windows_x86_64, Windows10 is better.
+ - Different pre-trained models are suitable for different effects.
+ - Run time depends on computer performance.
+ - If output video cannot be played, you can try with [potplayer](https://daumpotplayer.com/download/). ### Run from source #### Prerequisites -- Linux, Mac OS, Windows -- Python 3.6+ -- [ffmpeg 3.4](http://ffmpeg.org/) -- [Pytorch 1.0+](https://pytorch.org/) [(Old version codes)](https://github.com/HypoX64/DeepMosaics/tree/Pytorch0.4) -- CPU or NVIDIA GPU + CUDA CuDNN
+ - Linux, Mac OS, Windows + - Python 3.6+ + - [ffmpeg 3.4](http://ffmpeg.org/) + - [Pytorch 1.0+](https://pytorch.org/) [(Old version codes)](https://github.com/HypoX64/DeepMosaics/tree/Pytorch0.4) + - CPU or NVIDIA GPU + CUDA CuDNN
#### Dependencies This code depends on opencv-python, torchvision available via pip install. -#### Clone this repo: +#### Clone this repo ```bash git clone https://github.com/HypoX64/DeepMosaics cd DeepMosaics diff --git a/README_CN.md b/README_CN.md index dcd9537..b8ac9ed 100644 --- a/README_CN.md +++ b/README_CN.md @@ -16,21 +16,21 @@ ![image](./imgs/GUI.png)
注意事项:
-- 程序的运行要求在64位Windows操作系统,我仅在Windows10下运行过,其他版本暂未经过测试
-- 请根据需求选择合适的预训练模型进行测试
-- 运行时间取决于电脑性能,对于视频文件,我们建议可以先使用截图进行测试.
-- 如果输出的视频无法播放,这边建议您可以尝试[potplayer](https://daumpotplayer.com/download/). + - 程序的运行要求在64位Windows操作系统,我仅在Windows10运行过,其他版本暂未经过测试
+ - 请根据需求选择合适的预训练模型进行测试
+ - 运行时间取决于电脑性能,对于视频文件,我们建议可以先使用截图进行测试.
+ - 如果输出的视频无法播放,这边建议您尝试[potplayer](https://daumpotplayer.com/download/). ### 通过源代码运行 #### 前提要求 -- Linux, Mac OS, Windows -- Python 3.6+ -- [ffmpeg 3.4](http://ffmpeg.org/) -- [Pytorch 1.0+](https://pytorch.org/) [(Old version codes)](https://github.com/HypoX64/DeepMosaics/tree/Pytorch0.4) -- CPU or NVIDIA GPU + CUDA CuDNN
+ - Linux, Mac OS, Windows + - Python 3.6+ + - [ffmpeg 3.4](http://ffmpeg.org/) + - [Pytorch 1.0+](https://pytorch.org/) [(Old version codes)](https://github.com/HypoX64/DeepMosaics/tree/Pytorch0.4) + - CPU or NVIDIA GPU + CUDA CuDNN
#### Python依赖项 代码依赖于opencv-python以及 torchvision,可有通过pip install 进行安装. -#### 克隆源代码: +#### 克隆源代码 ```bash git clone https://github.com/HypoX64/DeepMosaics cd DeepMosaics diff --git a/deepmosaic.py b/deepmosaic.py index fc78e25..2787650 100644 --- a/deepmosaic.py +++ b/deepmosaic.py @@ -1,115 +1,120 @@ -import sys import os -import random - import numpy as np import cv2 -import torch 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) -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) - img = runmodel.add_mosaic_to_image(img,net,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) +def main(): + if opt.mode == 'add': - # 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) + 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() - 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')) + # 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) -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) + # 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) - 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() + 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')) - # get position - for imagepath in imagepaths: - imagepath=os.path.join('./tmp/video2image',imagepath) - img_origin = impro.imread(imagepath) + 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) - 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: + 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_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')) + 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) -util.clean_tempfiles(tmp_init = False) \ No newline at end of file +if __name__ == '__main__': + try: + main() + except Exception as e: + print('Error:',e) + input('Please press any key to exit.\n') + util.clean_tempfiles(tmp_init = False) + exit(0) diff --git a/util/mosaic.py b/util/mosaic.py index 99fbf92..54cc1a3 100755 --- a/util/mosaic.py +++ b/util/mosaic.py @@ -69,13 +69,13 @@ def addmosaic_autosize(img,mask,model,area_type = 'normal'): area = 0 area = area/(alpha*alpha) if area>50000: - img_mosaic = addmosaic_normal(img,mask,alpha*((area-50000)/50000+16),model = model) + img_mosaic = addmosaic_normal(img,mask,alpha*((area-50000)/50000+12),model = model) elif 20000