fix auto mosaic size

README.md

@@ -17,21 +17,21 @@ Download this version via [[Google Drive]](https://drive.google.com/open?id=1LTE
 ![image](./imgs/GUI.png)<br>
 
 Attentions:<br>
-- Require Windows_x86_64, Windows10 is better.<br>
-- Different pre-trained models are suitable for different effects.<br>
-- Run time depends on computer performance.<br>
-- If output video cannot be played, you can try with [potplayer](https://daumpotplayer.com/download/).
+  - Require Windows_x86_64, Windows10 is better.<br>
+  - Different pre-trained models are suitable for different effects.<br>
+  - Run time depends on computer performance.<br>
+  - 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<br>
+  - 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<br>
 #### 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

README_CN.md

@@ -16,21 +16,21 @@
 ![image](./imgs/GUI.png)<br>
 
 注意事项:<br>
-- 程序的运行要求在64位Windows操作系统,我仅在Windows10下运行过,其他版本暂未经过测试<br>
-- 请根据需求选择合适的预训练模型进行测试<br>
-- 运行时间取决于电脑性能,对于视频文件,我们建议可以先使用截图进行测试.<br>
-- 如果输出的视频无法播放,这边建议您可以尝试[potplayer](https://daumpotplayer.com/download/).
+  - 程序的运行要求在64位Windows操作系统,我仅在Windows10运行过,其他版本暂未经过测试<br>
+  - 请根据需求选择合适的预训练模型进行测试<br>
+  - 运行时间取决于电脑性能,对于视频文件,我们建议可以先使用截图进行测试.<br>
+  - 如果输出的视频无法播放,这边建议您尝试[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<br>
+  - 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<br>
 #### Python依赖项
 代码依赖于opencv-python以及 torchvision,可有通过pip install 进行安装.
-#### 克隆源代码:
+#### 克隆源代码
 ```bash
 git clone https://github.com/HypoX64/DeepMosaics
 cd DeepMosaics

deepmosaic.py

-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)

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<area<=50000:
-        img_mosaic = addmosaic_normal(img,mask,alpha*((area-20000)/30000+12),model = model)
+        img_mosaic = addmosaic_normal(img,mask,alpha*((area-20000)/30000+8),model = model)
     elif 5000<area<=20000:
-        img_mosaic = addmosaic_normal(img,mask,alpha*((area-5000)/20000+8),model = model)
+        img_mosaic = addmosaic_normal(img,mask,alpha*((area-5000)/20000+7),model = model)
     elif 0<=area<=5000:
-        img_mosaic = addmosaic_normal(img,mask,alpha*((area-0)/5000+4),model = model)
+        img_mosaic = addmosaic_normal(img,mask,alpha*((area-0)/5000+6),model = model)
     else:
         pass
     return img_mosaic