diff --git a/README.md b/README.md
index f85ccba22ce84f816405d777fb20766f296a865c..9e9822b15f1e8dcc6b10cb391f354b63d0b37477 100755
--- a/README.md
+++ b/README.md
@@ -2,6 +2,7 @@
# 
DeepMosaics
You can use it to automatically remove the mosaics in images and videos, or add mosaics to them.
This porject based on ‘semantic segmentation’ and ‘Image-to-Image Translation’.
+
* [中文版](./README_CN.md)
### More example
@@ -16,9 +17,6 @@ mosaic image | DeepCreamPy | ours
 |  | 
 |  | 
-## Notice
-The code do not include the part of training, I will finish it in my free time.
-
## Run DeepMosaics
You can either run DeepMosaics via pre-built binary package or from source.
@@ -39,7 +37,7 @@ Attentions:
#### Prerequisites
- Linux, Mac OS, Windows
- Python 3.6+
- - [ffmpeg 3.4](http://ffmpeg.org/)
+ - [ffmpeg 3.4.6](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
@@ -51,7 +49,7 @@ cd DeepMosaics
```
#### Get pre_trained models and test video
You can download pre_trained models and test video and replace the files in the project.
-[[Google Drive]](https://drive.google.com/open?id=10nARsiZoZGcaKw40nQu9fJuRp1oeabPs) [[百度云,提取码7thu]](https://pan.baidu.com/s/1IG4bdIiIC9PH9-oEyae5Sg)
+[[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs) [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ)
#### Simple example
* Add Mosaic (output video will save in './result')
@@ -64,8 +62,7 @@ python3 deepmosaic.py --mode clean --model_path ./pretrained_models/clean_hands_
```
#### More parameters
If you want to test other image or video, please refer to this file.
-[[options.py]](https://github.com/HypoX64/DeepMosaics/blob/master/options.py)
-
+[[options.py]](./cores/options.py)
## Acknowledgments
This code borrows heavily from [[pytorch-CycleGAN-and-pix2pix]](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix) [[Pytorch-UNet]](https://github.com/milesial/Pytorch-UNet)[[pix2pixHD]](https://github.com/NVIDIA/pix2pixHD).
diff --git a/README_CN.md b/README_CN.md
index 0797e29f52a887737dd5ea823cc3bbb611ef891c..aae9ca5ae92482c00b39dc2eac26f2baf641a4e1 100644
--- a/README_CN.md
+++ b/README_CN.md
@@ -15,8 +15,9 @@
 |  | 
## 一些说明
-代码暂不包含训练部分,训练方法我将在空闲时间给出.
-现在,代码已经支持基于[pix2pixHD](https://github.com/NVIDIA/pix2pixHD)训练出的模型,但网络仍在训练中,这将使得输出结果看起来更加清晰,"真实".
+* 训练部分并不完全
+* 现在,代码已经支持基于[pix2pixHD](https://github.com/NVIDIA/pix2pixHD)训练出的模型,但网络仍在训练中,这将使得输出结果看起来更加清晰,"真实".
+* 新的模型,可根据视频帧间关系进行马赛克恢复,在pretrained model 中被命名为*_video.pth
## 如何运行
可以通过我们预编译好的二进制包或源代码运行.
@@ -38,7 +39,7 @@
#### 前提要求
- Linux, Mac OS, Windows
- Python 3.6+
- - [ffmpeg 3.4](http://ffmpeg.org/)
+ - [ffmpeg 3.4.6](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依赖项
@@ -50,7 +51,7 @@ cd DeepMosaics
```
#### 下载测试视频以及预训练模型
可以通过以下两种方法下载测试视频以及预训练模型,并将他们置于项目文件夹中.
-[[Google Drive]](https://drive.google.com/open?id=10nARsiZoZGcaKw40nQu9fJuRp1oeabPs) [[百度云,提取码7thu]](https://pan.baidu.com/s/1IG4bdIiIC9PH9-oEyae5Sg)
+[[Google Drive]](https://drive.google.com/open?id=1LTERcN33McoiztYEwBxMuRjjgxh4DEPs) [[百度云,提取码1x0a]](https://pan.baidu.com/s/10rN3U3zd5TmfGpO_PEShqQ)
#### 简单的例子
* 为视频添加马赛克,例子中认为手是需要打码的区域 ,可以通过切换预训练模型切换自动打码区域(输出结果将储存到 './result')
@@ -63,7 +64,7 @@ python3 deepmosaic.py --mode clean --model_path ./pretrained_models/clean_hands_
```
#### 更多的参数
如果想要测试其他的图片或视频,请参照以下文件输入参数.
-[[options.py]](https://github.com/HypoX64/DeepMosaics/blob/master/options.py)
+[[options.py]](./cores/options.py)
## 鸣谢
代码大量的参考了以下项目:[[pytorch-CycleGAN-and-pix2pix]](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix) [[Pytorch-UNet]](https://github.com/milesial/Pytorch-UNet)[[pix2pixHD]](https://github.com/NVIDIA/pix2pixHD)..
diff --git a/cores/__init__.py b/cores/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..59321c09934652a995b18883b7d0406e66be1a5e
--- /dev/null
+++ b/cores/__init__.py
@@ -0,0 +1 @@
+from .options import *
\ No newline at end of file
diff --git a/cores/core.py b/cores/core.py
new file mode 100644
index 0000000000000000000000000000000000000000..c544a4063bdba2d179c0276b501c9dbfd6a159b6
--- /dev/null
+++ b/cores/core.py
@@ -0,0 +1,164 @@
+import os
+import numpy as np
+import cv2
+
+from models import runmodel,loadmodel
+from util import mosaic,util,ffmpeg,filt,data
+from util import image_processing as impro
+
+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 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'))
+
+def cleanmosaic_img(opt):
+ netG = loadmodel.pix2pix(opt)
+ net_mosaic_pos = loadmodel.unet_clean(opt)
+ path = opt.media_path
+ print('Clean Mosaic:',path)
+ img_origin = impro.imread(path)
+ x,y,size = runmodel.get_mosaic_position(img_origin,net_mosaic_pos,opt)[:3]
+ 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)
+ else:
+ print('Do not find mosaic')
+ cv2.imwrite(os.path.join(opt.result_dir,os.path.splitext(os.path.basename(path))[0]+'_clean.jpg'),img_result)
+
+def cleanmosaic_video_byframe(opt):
+ netG = loadmodel.pix2pix(opt)
+ net_mosaic_pos = loadmodel.unet_clean(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)
+ positions = []
+ imagepaths=os.listdir('./tmp/video2image')
+ imagepaths.sort()
+
+ # get position
+ for imagepath in imagepaths:
+ img_origin = impro.imread(os.path.join('./tmp/video2image',imagepath))
+ x,y,size = runmodel.get_mosaic_position(img_origin,net_mosaic_pos,opt)[:3]
+ 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):
+ 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('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'))
+
+def cleanmosaic_video_fusion(opt):
+ net = loadmodel.pix2pix(opt)
+ net_mosaic_pos = loadmodel.unet_clean(opt)
+ path = opt.media_path
+ N = 25
+
+ 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:
+ 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,net_mosaic_pos,opt)
+ cv2.imwrite(os.path.join('./tmp/mosaic_mask',imagepath), mask)
+ 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
+ print('Clean mosaic...')
+ for i,imagepath in enumerate(imagepaths,0):
+ print('Clean mosaic:',imagepath)
+ 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((256,256,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,256)
+ 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, 256)
+ mosaic_input[:,:,-1] = mask
+ mosaic_input = data.im2tensor(mosaic_input,bgr2rgb=False,use_gpu=opt.use_gpu,use_transform = False)
+ unmosaic_pred = net(mosaic_input)
+
+ unmosaic_pred = (unmosaic_pred.cpu().detach().numpy()*255)[0]
+ img_fake = unmosaic_pred.transpose((1, 2, 0))
+ 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)
+
+ 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'))
\ No newline at end of file
diff --git a/options.py b/cores/options.py
similarity index 86%
rename from options.py
rename to cores/options.py
index 786070229a80865831584d6067232203b8622651..a1bd9abae58ed70979880ffda2928f58fc2c19ac 100644
--- a/options.py
+++ b/cores/options.py
@@ -1,5 +1,6 @@
import argparse
import os
+import torch
class Options():
def __init__(self):
@@ -9,7 +10,8 @@ class Options():
def initialize(self):
#base
- self.parser.add_argument('--use_gpu', action='store_true', help='if input it, use gpu')
+ self.parser.add_argument('--use_gpu',type=bool,default=True, help='if True, use gpu')
+ # self.parser.add_argument('--use_gpu', action='store_true', help='if input it, use gpu')
self.parser.add_argument('--media_path', type=str, default='./hands_test.mp4',help='your videos or images path')
self.parser.add_argument('--mode', type=str, default='auto',help='add or clean mosaic into your media auto | add | clean')
self.parser.add_argument('--model_path', type=str, default='./pretrained_models/add_hands_128.pth',help='pretrained model path')
@@ -24,7 +26,7 @@ class Options():
self.parser.add_argument('--output_size', type=int, default=0,help='size of output file,if 0 -> origin')
#CleanMosaic
- self.parser.add_argument('--netG', type=str, default='auto',help='select model to use for netG(clean mosaic) -> auto | unet_128 | resnet_9blocks | HD')
+ self.parser.add_argument('--netG', type=str, default='auto',help='select model to use for netG(clean mosaic) -> auto | unet_128 | resnet_9blocks | HD | video')
self.parser.add_argument('--mosaic_position_model_path', type=str, default='auto',help='name of model use to find mosaic position')
self.parser.add_argument('--no_feather', action='store_true', help='if true, no edge feather and color correction, but run faster')
self.parser.add_argument('--medfilt_num', type=int, default=11,help='medfilt window of mosaic movement in the video')
@@ -36,6 +38,12 @@ class Options():
self.initialize()
self.opt = self.parser.parse_args()
+ if torch.cuda.is_available() and self.opt.use_gpu:
+ self.opt.use_gpu = True
+ else:
+ self.opt.use_gpu = False
+
+
if self.opt.mode == 'auto':
if 'add' in self.opt.model_path:
self.opt.mode = 'add'
@@ -51,6 +59,8 @@ class Options():
self.opt.netG = 'resnet_9blocks'
elif 'HD' in self.opt.model_path:
self.opt.netG = 'HD'
+ elif 'video' in self.opt.model_path:
+ self.opt.netG = 'video'
else:
print('Type of Generator error!')
diff --git a/deepmosaic.py b/deepmosaic.py
index 2c5bff765336baa0c7c36edf6ae1e0cc1b269ee3..82bde84e6b75bbcc6ab71d8defc579ffe4d29e53 100644
--- a/deepmosaic.py
+++ b/deepmosaic.py
@@ -1,113 +1,31 @@
import os
-import numpy as np
-import cv2
-
-from models import runmodel,loadmodel
-from util import mosaic,util,ffmpeg,filt
-from util import image_processing as impro
-from options import Options
+from cores import Options,core
+from util import util
opt = Options().getparse()
util.file_init(opt)
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'))
+ 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':
- 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)
+ 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')
- # 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)
if __name__ == '__main__':
diff --git a/models/loadmodel.py b/models/loadmodel.py
index 0dd65ced085148553eab6d81565e39e555b565b9..8126ee89963a3152af7cceeae4c928fe34a7c893 100755
--- a/models/loadmodel.py
+++ b/models/loadmodel.py
@@ -7,6 +7,8 @@ def pix2pix(opt):
# print(opt.model_path,opt.netG)
if opt.netG == 'HD':
netG = define_G_HD(3, 3, 64, 'global' ,4)
+ elif opt.netG == 'video':
+ netG = define_G(3*25+1, 3, 128, 'unet_128', norm='instance',use_dropout=True, init_type='normal', gpu_ids=[])
else:
netG = define_G(3, 3, 64, opt.netG, norm='batch',use_dropout=True, init_type='normal', gpu_ids=[])
diff --git a/models/runmodel.py b/models/runmodel.py
index b41e1cd97a3f927074c3b735aec59a4a0ece517c..31a4f8c453784755cc7fe2977d08166da16588d8 100755
--- a/models/runmodel.py
+++ b/models/runmodel.py
@@ -40,10 +40,10 @@ def get_ROI_position(img,net,opt):
x,y,halfsize,area = impro.boundingSquare(mask, 1)
return mask,x,y,area
-def get_mosaic_position(img_origin,net_mosaic_pos,opt):
+def get_mosaic_position(img_origin,net_mosaic_pos,opt,threshold = 128 ):
mask = run_unet_rectim(img_origin,net_mosaic_pos,use_gpu = opt.use_gpu)
- mask = impro.mask_threshold(mask,10,128)
+ mask = impro.mask_threshold(mask,10,threshold)
x,y,size,area = impro.boundingSquare(mask,Ex_mul=1.5)
rat = min(img_origin.shape[:2])/128.0
x,y,size = int(rat*x),int(rat*y),int(rat*size)
- return x,y,size
\ No newline at end of file
+ return x,y,size,mask
\ No newline at end of file
diff --git a/util/data.py b/util/data.py
index 1978c96412387a523f59665c761077a0545ddc2e..2396fce365cc9be811e44d01e68b92edb3143a19 100755
--- a/util/data.py
+++ b/util/data.py
@@ -19,13 +19,28 @@ def tensor2im(image_tensor, imtype=np.uint8, rgb2bgr = True):
return image_numpy.astype(imtype)
-def im2tensor(image_numpy, imtype=np.uint8, bgr2rgb = True, reshape = True, use_gpu = True):
- h, w = image_numpy.shape[:2]
+def im2tensor(image_numpy, imtype=np.uint8, bgr2rgb = True, reshape = True, use_gpu = True, use_transform = True):
+ h, w ,ch = image_numpy.shape
if bgr2rgb:
image_numpy = image_numpy[...,::-1]-np.zeros_like(image_numpy)
- image_tensor = transform(image_numpy)
+ if use_transform:
+ image_tensor = transform(image_numpy)
+ else:
+ image_numpy = image_numpy/255.0
+ image_numpy = image_numpy.transpose((2, 0, 1))
+ image_tensor = torch.from_numpy(image_numpy).float()
if reshape:
- image_tensor=image_tensor.reshape(1,3,h,w)
+ image_tensor=image_tensor.reshape(1,ch,h,w)
if use_gpu:
image_tensor = image_tensor.cuda()
- return image_tensor
\ No newline at end of file
+ return image_tensor
+
+# def im2tensor(image_numpy, use_gpu=False):
+# h, w ,ch = image_numpy.shape
+# image_numpy = image_numpy/255.0
+# image_numpy = image_numpy.transpose((2, 0, 1))
+# image_numpy = image_numpy.reshape(-1,ch,h,w)
+# img_tensor = torch.from_numpy(image_numpy).float()
+# if use_gpu:
+# img_tensor = img_tensor.cuda()
+# return img_tensor
\ No newline at end of file
diff --git a/util/mosaic.py b/util/mosaic.py
index 54cc1a3ecddeaec54fe5fb6c2f6e03bd03f9868a..e2ba16de3a32fad6ccddb986b99d85d9d437b11e 100755
--- a/util/mosaic.py
+++ b/util/mosaic.py
@@ -13,7 +13,7 @@ def addmosaic(img,mask,opt):
img = addmosaic_normal(img,mask,opt.mosaic_size,opt.output_size,model = opt.mosaic_mod)
return img
-def addmosaic_normal(img,mask,n,out_size = 0,model = 'squa_avg'):
+def addmosaic_normal(img,mask,n,out_size = 0,model = 'squa_avg',rect_rat = 1.6):
n = int(n)
if out_size:
img = resize(img,out_size)
@@ -27,6 +27,12 @@ def addmosaic_normal(img,mask,n,out_size = 0,model = 'squa_avg'):
if mask[int(i*n+n/2),int(j*n+n/2)] == 255:
img_mosaic[i*n:(i+1)*n,j*n:(j+1)*n,:]=img[i*n:(i+1)*n,j*n:(j+1)*n,:].mean(0).mean(0)
+ elif model=='squa_mid':
+ for i in range(int(h/n)):
+ for j in range(int(w/n)):
+ if mask[int(i*n+n/2),int(j*n+n/2)] == 255:
+ img_mosaic[i*n:(i+1)*n,j*n:(j+1)*n,:]=img[i*n+int(n/2),j*n+int(n/2),:]
+
elif model == 'squa_random':
for i in range(int(h/n)):
for j in range(int(w/n)):
@@ -45,9 +51,8 @@ def addmosaic_normal(img,mask,n,out_size = 0,model = 'squa_avg'):
img_mosaic = cv2.add(imgroi1,imgroi2)
elif model =='rect_avg':
- rect_ratio=1+0.6*random.random()
n_h=n
- n_w=int(n*rect_ratio)
+ n_w=int(n*rect_rat)
for i in range(int(h/n_h)):
for j in range(int(w/n_w)):
if mask[int(i*n_h+n_h/2),int(j*n_w+n_w/2)] == 255:
@@ -55,6 +60,31 @@ def addmosaic_normal(img,mask,n,out_size = 0,model = 'squa_avg'):
return img_mosaic
+def get_autosize(img,mask,area_type = 'normal'):
+ h,w = img.shape[:2]
+ mask = cv2.resize(mask,(w,h))
+ alpha = np.min((w,h))/512
+ try:
+ if area_type == 'normal':
+ area = mask_area(mask)
+ elif area_type == 'bounding':
+ w,h = cv2.boundingRect(mask)[2:]
+ area = w*h
+ except:
+ area = 0
+ area = area/(alpha*alpha)
+ if area>50000:
+ size = alpha*((area-50000)/50000+12)
+ elif 20000