support new model

36383527 · hypox64 · d09e2ce5 · 36383527 · 36383527 · 36383527
11 changed file
--- a/README.md
+++ b/README.md
@@ -2,6 +2,7 @@
 # <img src="./imgs/icon.jpg" width="48">DeepMosaics
 You can use it to automatically remove the mosaics in images and videos, or add mosaics to them.<br>
 This porject based on ‘semantic segmentation’ and ‘Image-to-Image Translation’.<br>
 * [中文版](./README_CN.md)<br>
 ### More example
@@ -16,9 +17,6 @@ mosaic image | DeepCreamPy | ours
 ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/face_a_mosaic.jpg) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/a_dcp.png) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/face_a_clean.jpg) 
 ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/face_b_mosaic.jpg) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/b_dcp.png) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/face_b_clean.jpg) 
-## Notice
-The code do not include the part of training, I will finish it in my free time.<br>
 ## Run DeepMosaics
 You can either run DeepMosaics via pre-built binary package or from source.<br>
@@ -39,7 +37,7 @@ Attentions:<br>
 #### 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<br>
 #### 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.<br>
-[[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) <br>
-<br>
 ## 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).
--- a/README_CN.md
+++ b/README_CN.md
@@ -15,8 +15,9 @@
 ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/face_b_mosaic.jpg) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/b_dcp.png) | ![image](https://github.com/HypoX64/DeepMosaics_example/blob/master/face_b_clean.jpg) 
 ## 一些说明
-代码暂不包含训练部分,训练方法我将在空闲时间给出.<br>
+* 训练部分并不完全<br>
-现在,代码已经支持基于[pix2pixHD](https://github.com/NVIDIA/pix2pixHD)训练出的模型,但网络仍在训练中,这将使得输出结果看起来更加清晰,"真实".<br>
+* 现在,代码已经支持基于[pix2pixHD](https://github.com/NVIDIA/pix2pixHD)训练出的模型,但网络仍在训练中,这将使得输出结果看起来更加清晰,"真实".<br>
+* 新的模型,可根据视频帧间关系进行马赛克恢复,在pretrained model 中被命名为*_video.pth<br>
 ## 如何运行
 可以通过我们预编译好的二进制包或源代码运行.<br>
@@ -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<br>
 #### Python依赖项
@@ -50,7 +51,7 @@ cd DeepMosaics
 ```
 #### 下载测试视频以及预训练模型
 可以通过以下两种方法下载测试视频以及预训练模型,并将他们置于项目文件夹中.<br>
-[[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) <br>
 #### 简单的例子
 * 为视频添加马赛克,例子中认为手是需要打码的区域 ,可以通过切换预训练模型切换自动打码区域(输出结果将储存到 './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) <br>
+[[options.py]](./cores/options.py) <br>
 ## 鸣谢
 代码大量的参考了以下项目:[[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)..
--- a/cores/__init__.py
+++ b/cores/__init__.py
+from .options import *
\ No newline at end of file
--- a/cores/core.py
+++ b/cores/core.py
+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
--- a/options.py
+++ b/options.py
 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!')

--- a/deepmosaic.py
+++ b/deepmosaic.py
 import os
-import numpy as np
+from cores import Options,core
-import cv2
+from util import util
-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)
 def main():
    if opt.mode == 'add':
+        if util.is_img(opt.media_path):
-        net = loadmodel.unet(opt)
+            core.addmosaic_img(opt)
-        path = opt.media_path
+        elif util.is_video(opt.media_path):
-        if util.is_img(path):
+            core.addmosaic_video(opt)
-            print('Add Mosaic:',path)
+        else:
-            img = impro.imread(path)
+            print('This type of file is not supported')
-            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)
+        if util.is_img(opt.media_path):
-        net_mosaic_pos = loadmodel.unet_clean(opt)
+            core.cleanmosaic_img(opt)
-        path = opt.media_path
+        elif util.is_video(opt.media_path):
-        if util.is_img(path):
+            if opt.netG == 'video':
-            print('Clean Mosaic:',path)
+                core.cleanmosaic_video_fusion(opt)
-            img_origin = impro.imread(path)
+            else:
-            x,y,size = runmodel.get_mosaic_position(img_origin,net_mosaic_pos,opt)
+                core.cleanmosaic_video_byframe(opt)
-            img_result = img_origin.copy()
+        else:
-            if size != 0 :
+            print('This type of file is not supported')
-                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)
 if __name__ == '__main__':

--- 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=[])

--- 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
+    return x,y,size,mask
\ No newline at end of file
--- 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):
+def im2tensor(image_numpy, imtype=np.uint8, bgr2rgb = True, reshape = True, use_gpu = True,  use_transform = True):
-    h, w = image_numpy.shape[:2]
+    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
+# 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
--- 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<area<=50000:
+        size = alpha*((area-20000)/30000+8)
+    elif 5000<area<=20000:
+        size = alpha*((area-5000)/20000+7)
+    elif 0<=area<=5000:
+        size = alpha*((area-0)/5000+6)
+    else:
+        pass
+    return size
 def addmosaic_autosize(img,mask,model,area_type = 'normal'):
    h,w = img.shape[:2]
    mask = cv2.resize(mask,(w,h))
@@ -109,11 +139,11 @@ def addmosaic_random(img,mask,area_type = 'normal'):
 def random_mod(img,mask,n):
    ran=random.random()
-    if ran < 0.1:
+    if ran < 0.2:
-        img = addmosaic_normal(img,mask,n,model = 'squa_random')
+        img = addmosaic_normal(img,mask,n,model = 'squa_mid')
-    if 0.1 <= ran < 0.3:
+    if 0.2 <= ran < 0.4:
        img = addmosaic_normal(img,mask,n,model = 'squa_avg')
-    elif 0.3 <= ran <0.5:
+    elif 0.4 <= ran <0.6:
        img = addmosaic_normal(img,mask,n,model = 'squa_avg_circle_edge')
    else:
        img = addmosaic_normal(img,mask,n,model = 'rect_avg')

--- a/util/util.py
+++ b/util/util.py
@@ -29,6 +29,7 @@ def is_video(path):
 def  writelog(path,log):
    f = open(path,'a+')
    f.write(log+'\n')
+    f.close()
 def makedirs(path):
    if os.path.isdir(path):
@@ -46,7 +47,9 @@ def clean_tempfiles(tmp_init=True):
        os.makedirs('./tmp/addmosaic_image')
        os.makedirs('./tmp/mosaic_crop')
        os.makedirs('./tmp/replace_mosaic')
+        os.makedirs('./tmp/mosaic_mask')
        os.makedirs('./tmp/ROI_mask')
+        os.makedirs('./tmp/ROI_mask_check')
 def file_init(opt):
    if not os.path.isdir(opt.result_dir):