Create branch dev

.gitignore

@@ -141,6 +141,7 @@ dataset/
 test*
 video_tmp/
 result/
+nohup.out
 #./
 /pix2pix
 /pix2pixHD

README.md

@@ -67,11 +67,11 @@ You can download pre_trained models and put them into './pretrained_models'.<br>
 #### Simple example
 * Add Mosaic (output media will save in './result')<br>
 ```bash
-python3 deepmosaic.py --media_path ./imgs/ruoruo.jpg --model_path ./pretrained_models/mosaic/add_face.pth --use_gpu -1
+python3 deepmosaic.py --media_path ./imgs/ruoruo.jpg --model_path ./pretrained_models/mosaic/add_face.pth --use_gpu 0
 ```
 * Clean Mosaic (output media will save in './result')<br>
 ```bash
-python3 deepmosaic.py --media_path ./result/ruoruo_add.jpg --model_path ./pretrained_models/mosaic/clean_face_HD.pth --use_gpu -1
+python3 deepmosaic.py --media_path ./result/ruoruo_add.jpg --model_path ./pretrained_models/mosaic/clean_face_HD.pth --use_gpu 0
 ```
 #### More parameters
 If you want to test other image or video, please refer to this file.<br>

README_CN.md

@@ -64,11 +64,11 @@ cd DeepMosaics
 #### 简单的例子
 * 为视频添加马赛克,例子中认为脸是需要打码的区域 ,可以通过切换预训练模型切换自动打码区域(输出结果将储存到 './result')<br>
 ```bash
-python3 deepmosaic.py --media_path ./imgs/ruoruo.jpg --model_path ./pretrained_models/mosaic/add_face.pth --use_gpu -1
+python3 deepmosaic.py --media_path ./imgs/ruoruo.jpg --model_path ./pretrained_models/mosaic/add_face.pth --use_gpu 0
 ```
 * 将视频中的马赛克移除,对于不同的打码物体需要使用对应的预训练模型进行马赛克消除(输出结果将储存到  './result')<br>
 ```bash
-python3 deepmosaic.py --media_path ./result/ruoruo_add.jpg --model_path ./pretrained_models/mosaic/clean_face_HD.pth --use_gpu -1
+python3 deepmosaic.py --media_path ./result/ruoruo_add.jpg --model_path ./pretrained_models/mosaic/clean_face_HD.pth --use_gpu 0
 ```
 #### 更多的参数
 如果想要测试其他的图片或视频,请参照以下文件输入参数.<br>

cores/options.py

@@ -1,6 +1,6 @@
 import argparse
 import os
-import torch
+
 
 class Options():
     def __init__(self):
@@ -10,7 +10,7 @@ def __init__(self):
     def initialize(self):
 
         #base
-        self.parser.add_argument('--use_gpu',type=int,default=0, help='if -1, do not use gpu')
+        self.parser.add_argument('--use_gpu',type=int,default=0, help='if -1, use cpu')
         # self.parser.add_argument('--use_gpu', action='store_true', help='if input it, use gpu')
         self.parser.add_argument('--media_path', type=str, default='./imgs/ruoruo.jpg',help='your videos or images path')
         self.parser.add_argument('--mode', type=str, default='auto',help='Program running mode. auto | add | clean | style')
@@ -54,10 +54,12 @@ def getparse(self):
 
         model_name = os.path.basename(self.opt.model_path)
 
+        os.environ["CUDA_VISIBLE_DEVICES"] = str(self.opt.use_gpu)
+        import torch
         if torch.cuda.is_available() and self.opt.use_gpu > -1:
-            self.opt.use_gpu = True
+            pass
         else:
-            self.opt.use_gpu = False
+            self.opt.use_gpu = -1
 
         if self.opt.mode == 'auto':
             if 'clean' in model_name or self.opt.traditional:

deepmosaic.py

@@ -72,6 +72,6 @@ def main():
         for stack in traceback.extract_tb(ex_stack):
             print(stack)
         input('Please press any key to exit.\n')
-        util.clean_tempfiles(tmp_init = False)
+        #util.clean_tempfiles(tmp_init = False)
         exit(0)
 

docs/training_with_your_own_dataset.md

@@ -53,7 +53,7 @@ python make_video_dataset.py --datadir 'dir for your videos' --model_path ../pre
 ### Add
 ```bash
 cd train/add
-python train.py --gpu_id 0 --dataset ../../datasets/draw/face --savename face --loadsize 512 --finesize 360 --batchsize 16
+python train.py --use_gpu 0 --dataset ../../datasets/draw/face --savename face --loadsize 512 --finesize 360 --batchsize 16
 ```
 ### Clean
 * For image datasets(generated by ```make_pix2pix_dataset.py```)
@@ -67,7 +67,7 @@ python train.py --name face --resize_or_crop resize_and_crop --loadSize 563 --fi
 * For video datasets(generated by ```make_video_dataset.py```)
 ```bash
 cd train/clean
-python train.py --dataset ../../datasets/video/face --savename face --savefreq 100000 --gan --hd --lr 0.0002 --lambda_gan 1 --gpu_id 0 --perload_num 8
+python train.py --dataset ../../datasets/video/face --savename face --savefreq 100000 --gan --hd --lr 0.0002 --lambda_gan 1 --use_gpu 0
 ```
 ## Testing
 Put saved network to ```./pretrained_models/mosaic/``` and rename it as ```add_face.pth``` or ```clean_face_HD.pth``` or ```clean_face_video_HD.pth```

make_datasets/make_pix2pix_dataset.py

@@ -1,24 +1,24 @@
 import os
-import random
 import sys
+sys.path.append("..")
+from cores import Options
+opt = Options()
+
+import random
 import datetime
 import time
-import shutil
-import threading
 import warnings
 warnings.filterwarnings(action='ignore')
 
 import numpy as np
 import cv2
+import torch
 
-sys.path.append("..")
 from models import runmodel,loadmodel
 import util.image_processing as impro
 from util import util,mosaic,data
-from cores import Options
 
 
-opt = Options()
 opt.parser.add_argument('--datadir',type=str,default='../datasets/draw/face', help='')
 opt.parser.add_argument('--savedir',type=str,default='../datasets/pix2pix/face', help='')
 opt.parser.add_argument('--name',type=str,default='', help='save name')
@@ -61,51 +61,16 @@
     maskpaths.sort()
 if 'network' in opt.mod or 'irregular' in opt.mod:
     imgpaths = util.Traversal(opt.datadir)
+    imgpaths = util.is_imgs(imgpaths)
     random.shuffle (imgpaths)
 if 'irregular' in opt.mod:
     irrpaths = util.Traversal(opt.irrholedir)
 
 
-#def network
+#def network                
 if 'network' in opt.mod:
     net = loadmodel.bisenet(opt,'roi')
 
-
-# def checksaveimage(opt,img,mask):
-
-#     #check
-#     saveflag = True
-#     x,y,size,area = impro.boundingSquare(mask, random.uniform(1.4,1.6))
-#     if area < 1000:
-#         saveflag = False
-#     else:
-#         if opt.square:
-#             if size < opt.minsize:
-#                 saveflag = False
-#             else:
-#                 img = impro.resize(img[y-size:y+size,x-size:x+size],opt.outsize,interpolation=cv2.INTER_CUBIC)
-#                 mask =  impro.resize(mask[y-size:y+size,x-size:x+size],opt.outsize,interpolation=cv2.INTER_CUBIC)
-#                 if impro.Q_lapulase(img)<opt.quality:
-#                     saveflag = False         
-#         else:
-#             img = impro.resize(img,opt.outsize,interpolation=cv2.INTER_CUBIC)
-
-#     if saveflag:
-#         # add mosaic
-#         img_mosaic = mosaic.addmosaic_random(img, mask)
-#         global savecnt
-#         savecnt += 1
-
-#         if opt.hd:
-#             cv2.imwrite(os.path.join(train_A_path,opt.name+'%06d' % savecnt+'.jpg'), img_mosaic)
-#             cv2.imwrite(os.path.join(train_B_path,opt.name+'%06d' % savecnt+'.jpg'), img)
-#         else:
-#             merge_img = impro.makedataset(img_mosaic, img)
-#             cv2.imwrite(os.path.join(train_path,opt.name+'%06d' % savecnt+'.jpg'), merge_img)
-#         if opt.savemask:
-#             cv2.imwrite(os.path.join(mask_save_path,opt.name+'%06d' % savecnt+'.png'), mask)
-
-
 print('Find images:',len(imgpaths))
 starttime = datetime.datetime.now()
 filecnt = 0
@@ -126,6 +91,8 @@
                 mask = mask_irr
             if 'network' in opt.mod:
                 mask_net = runmodel.get_ROI_position(img,net,opt,keepsize=True)[0]
+                if opt.use_gpu != -1:
+                    torch.cuda.empty_cache()
                 if not opt.all_mosaic_area:
                     mask_net = impro.find_mostlikely_ROI(mask_net)
                 mask = mask_net
@@ -157,7 +124,12 @@
                 if saveflag:
                     # add mosaic
                     img_mosaic = mosaic.addmosaic_random(img, mask)
-                    # global savecnt
+                    # random blur
+                    if random.random()>0.5:
+                        Q = random.randint(1,15)
+                        img = impro.dctblur(img,Q)
+                        img_mosaic = impro.dctblur(img_mosaic,Q)
+
                     savecnt += 1
 
                     if opt.hd:
@@ -185,6 +157,6 @@
             all_time = used_time/filecnt*all_length
 
             print('\r','',str(filecnt)+'/'+str(all_length)+' ',
-                util.get_bar(percent,30),'',
+                util.get_bar(percent,25),'',
                 util.second2stamp(used_time)+'/'+util.second2stamp(all_time),
                 'f:'+str(savecnt),end= " ")

make_datasets/make_video_dataset.py

@@ -1,21 +1,22 @@
 import os
-import random
 import sys
+sys.path.append("..")
+from cores import Options
+opt = Options()
+
+import random
 import datetime
 import time
-import shutil
-import threading
 
 import numpy as np
 import cv2
+import torch
 
-sys.path.append("..")
 from models import runmodel,loadmodel
 import util.image_processing as impro
 from util import util,mosaic,data,ffmpeg
-from cores import Options
 
-opt = Options()
+
 opt.parser.add_argument('--datadir',type=str,default='your video dir', help='')
 opt.parser.add_argument('--savedir',type=str,default='../datasets/video/face', help='')
 opt.parser.add_argument('--interval',type=int,default=30, help='interval of split video ')
@@ -25,6 +26,7 @@
 opt.parser.add_argument('--outsize', type=int ,default= 286,help='')
 opt.parser.add_argument('--startcnt', type=int ,default= 0,help='')
 opt.parser.add_argument('--minsize', type=int ,default= 96,help='minimal roi size')
+opt.parser.add_argument('--no_sclectscene', action='store_true', help='')  
 opt = opt.getparse()
 
 
@@ -44,52 +46,31 @@
 starttime = datetime.datetime.now()
 for videopath in videopaths:
     try:
-        timestamps=[]
-        fps,endtime,height,width = ffmpeg.get_video_infos(videopath)
-        for cut_point in range(1,int((endtime-opt.time)/opt.interval)):
-            util.clean_tempfiles()
-            ffmpeg.video2image(videopath, './tmp/video2image/%05d.'+opt.tempimage_type,fps=1,
-                start_time = util.second2stamp(cut_point*opt.interval),last_time = util.second2stamp(opt.time))
-            imagepaths = util.Traversal('./tmp/video2image')
-            cnt = 0 
-            for i in range(opt.time):
-                img = impro.imread(imagepaths[i])
-                mask = runmodel.get_ROI_position(img,net,opt,keepsize=True)[0]
-                if not opt.all_mosaic_area:
-                    mask = impro.find_mostlikely_ROI(mask)
-                x,y,size,area = impro.boundingSquare(mask,Ex_mul=1)
-                if area > opt.minmaskarea and size>opt.minsize and impro.Q_lapulase(img)>opt.quality:
-                    cnt +=1
-            if cnt == opt.time:
-                # print(second)
-                timestamps.append(util.second2stamp(cut_point*opt.interval))
+        if opt.no_sclectscene:
+            timestamps=['00:00:00']
+        else:
+            timestamps=[]
+            fps,endtime,height,width = ffmpeg.get_video_infos(videopath)
+            for cut_point in range(1,int((endtime-opt.time)/opt.interval)):
+                util.clean_tempfiles()
+                ffmpeg.video2image(videopath, './tmp/video2image/%05d.'+opt.tempimage_type,fps=1,
+                    start_time = util.second2stamp(cut_point*opt.interval),last_time = util.second2stamp(opt.time))
+                imagepaths = util.Traversal('./tmp/video2image')
+                cnt = 0 
+                for i in range(opt.time):
+                    img = impro.imread(imagepaths[i])
+                    mask = runmodel.get_ROI_position(img,net,opt,keepsize=True)[0]
+                    if not opt.all_mosaic_area:
+                        mask = impro.find_mostlikely_ROI(mask)
+                    x,y,size,area = impro.boundingSquare(mask,Ex_mul=1)
+                    if area > opt.minmaskarea and size>opt.minsize and impro.Q_lapulase(img)>opt.quality:
+                        cnt +=1
+                if cnt == opt.time:
+                    # print(second)
+                    timestamps.append(util.second2stamp(cut_point*opt.interval))
         util.writelog(os.path.join(opt.savedir,'opt.txt'),videopath+'\n'+str(timestamps))
         #print(timestamps)
 
-        # util.clean_tempfiles()
-        # fps,endtime,height,width = ffmpeg.get_video_infos(videopath)
-        # # print(fps,endtime,height,width)
-        # ffmpeg.continuous_screenshot(videopath, './tmp/video2image', 1)
-
-        # # find where to cut
-        # print('Find where to cut...')
-        # timestamps=[]
-        # imagepaths = util.Traversal('./tmp/video2image')
-        # for second in range(int(endtime)):
-        #     if second%opt.interval==0:
-        #         cnt = 0 
-        #         for i in range(opt.time):
-        #             img = impro.imread(imagepaths[second+i])
-        #             mask = runmodel.get_ROI_position(img,net,opt)[0]
-        #             if not opt.all_mosaic_area:
-        #                 mask = impro.find_mostlikely_ROI(mask)
-        #             if impro.mask_area(mask) > opt.minmaskarea and impro.Q_lapulase(img)>opt.quality:
-        #                 # print(impro.mask_area(mask))
-        #                 cnt +=1
-        #         if cnt == opt.time:
-        #             # print(second)
-        #             timestamps.append(util.second2stamp(second))
-
         #generate datasets
         print('Generate datasets...')
         for timestamp in timestamps:
@@ -143,3 +124,5 @@
         util.writelog(os.path.join(opt.savedir,'opt.txt'), 
               videopath+'\n'+str(result_cnt)+'\n'+str(e))
     video_cnt +=1
+    if opt.use_gpu != -1:
+        torch.cuda.empty_cache()

models/loadmodel.py

@@ -34,7 +34,7 @@ def pix2pix(opt):
     show_paramsnumber(netG,'netG')
     netG.load_state_dict(torch.load(opt.model_path))
     netG.eval()
-    if opt.use_gpu:
+    if opt.use_gpu != -1:
         netG.cuda()
     return netG
 
@@ -60,7 +60,7 @@ def style(opt):
         __patch_instance_norm_state_dict(state_dict, netG, key.split('.'))
     netG.load_state_dict(state_dict)
 
-    if opt.use_gpu:
+    if opt.use_gpu != -1:
         netG.cuda()
     return netG
 
@@ -72,7 +72,7 @@ def video(opt):
     show_paramsnumber(netG,'netG')
     netG.load_state_dict(torch.load(opt.model_path))
     netG.eval()
-    if opt.use_gpu:
+    if opt.use_gpu != -1:
         netG.cuda()
     return netG
 
@@ -87,24 +87,6 @@ def bisenet(opt,type='roi'):
     elif type == 'mosaic':
         net.load_state_dict(torch.load(opt.mosaic_position_model_path))
     net.eval()
-    if opt.use_gpu:
+    if opt.use_gpu != -1:
         net.cuda()
     return net
-
-# def unet_clean(opt):
-#     net = UNet(n_channels = 3, n_classes = 1)
-#     show_paramsnumber(net,'segment')
-#     net.load_state_dict(torch.load(opt.mosaic_position_model_path))
-#     net.eval()
-#     if opt.use_gpu:
-#         net.cuda()
-#     return net
-
-# def unet(opt):
-#     net = UNet(n_channels = 3, n_classes = 1)
-#     show_paramsnumber(net,'segment')
-#     net.load_state_dict(torch.load(opt.model_path))
-#     net.eval()
-#     if opt.use_gpu:
-#         net.cuda()
-#     return net