make more training data

.gitignore

@@ -141,6 +141,7 @@ test*/
 video_tmp/
 result/
 #./
+/pix2pixHD
 /tmp
 /to_make_show
 /test_media

README.md

@@ -2,7 +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>
-
+Master is not stable. Please use a [stable version](https://github.com/HypoX64/DeepMosaics/tree/stable)<br>
 * [中文版](./README_CN.md)<br>
 
 ### More example

README_CN.md

@@ -1,7 +1,7 @@
 ![image](./imgs/hand.gif)
 # <img src="./imgs/icon.jpg" width="48">DeepMosaics
 这是一个通过深度学习自动的为图片/视频添加马赛克,或消除马赛克的项目.<br>它基于“语义分割”以及“图像翻译”.<br>
-
+主分支并不稳定,请移步[稳定版本](https://github.com/HypoX64/DeepMosaics/tree/stable)<br>
 ### 更多例子
 原始 | 自动打码 |  自动去码  
 :-:|:-:|:-:

make_datasets/get_image_from_video.py

@@ -0,0 +1,18 @@
+import os
+import numpy as np
+import cv2
+import random
+import csv
+
+import sys
+sys.path.append("..")
+from util import util,ffmpeg
+from util import image_processing as impro
+
+files = util.Traversal('/media/hypo/Media/download')
+videos = util.is_videos(files)
+output_dir = './dataset/v2im'
+FPS = 1
+util.makedirs(output_dir)
+for video in videos:
+    ffmpeg.continuous_screenshot(video, output_dir, FPS)

make_datasets/use_addmosaic_model_make_video_dataset.py

@@ -22,7 +22,7 @@
 suffix = ''
 
 net = loadmodel.unet(opt)
-for path in videos:
+for i,path in enumerate(videos,0):
     try:
         path = os.path.join('./video',path)
         util.clean_tempfiles()
@@ -37,14 +37,14 @@
         mask_avg = np.zeros((impro.resize(img_ori_example, 128)).shape[:2])
         for imagepath in imagepaths:
             imagepath = os.path.join('./tmp/video2image',imagepath)
-            print('Find ROI location:',imagepath)
+            #print('Find ROI location:',imagepath)
             img = impro.imread(imagepath)
-            x,y,size,mask = runmodel.get_mosaic_position(img,net,opt,threshold = 64)
+            x,y,size,mask = runmodel.get_mosaic_position(img,net,opt,threshold = 80)
             cv2.imwrite(os.path.join('./tmp/ROI_mask',
                               os.path.basename(imagepath)),mask)
             positions.append([x,y,size])
             mask_avg = mask_avg + mask
-        print('Optimize ROI locations...')
+        #print('Optimize ROI locations...')
         mask_index = filt.position_medfilt(np.array(positions), 13)
 
         mask = np.clip(mask_avg/len(imagepaths),0,255).astype('uint8')
@@ -62,7 +62,7 @@
             os.makedirs(mask_path)
             os.makedirs(ori_path)
             os.makedirs(mosaic_path)
-            print('Add mosaic to images...')
+            #print('Add mosaic to images...')
             mosaic_size = mosaic.get_autosize(img_ori_example,mask,area_type = Area_Type)*random.uniform(1,2)
             models = ['squa_avg','rect_avg','squa_mid']
             mosaic_type = random.randint(0,len(models)-1)
@@ -81,4 +81,6 @@
                 cv2.imwrite(os.path.join(mosaic_path,os.path.basename(imagepaths[i])),img_mosaic_crop)
                 cv2.imwrite(os.path.join(mask_path,os.path.basename(imagepaths[i])),mask_crop)
     except Exception as e:
-        print(e)
+        print(e)
+
+    print(util.get_bar(100*i/len(videos),num=50))

make_datasets/use_irregular_holes_mask_make_dataset.py

@@ -11,19 +11,20 @@
 import datetime
 
 ir_mask_path = './Irregular_Holes_mask'
-img_path ='/home/hypo/桌面/FaceRankSample' 
-output_dir = './datasets_img'
-util.makedirs(output_dir)
+img_dir ='/home/hypo/MyProject/Haystack/CV/output/all/face' 
 MOD = 'HD' #HD | pix2pix | mosaic
 MASK = False # if True, output mask,too
-BOUNDING = False # if true the mosaic size will be more big
+BOUNDING = True # if true the mosaic size will be more big
+suffix = ''
+output_dir = os.path.join('./dataset_img',MOD)
+util.makedirs(output_dir)
 
-if MOD='HD':
+if MOD == 'HD':
     train_A_path = os.path.join(output_dir,'train_A')
     train_B_path = os.path.join(output_dir,'train_B')
     util.makedirs(train_A_path)
     util.makedirs(train_B_path)
-elif MOD='pix2pix':
+elif MOD == 'pix2pix':
     train_path = os.path.join(output_dir,'train')
     util.makedirs(train_path)
 if MASK:
@@ -42,12 +43,12 @@
  ])
 
 mask_names = os.listdir(ir_mask_path)
-img_names = os.listdir(img_path)
+img_names = os.listdir(img_dir)
 print('Find images:',len(img_names))
 
 for i,img_name in enumerate(img_names,1):
-    try:
-        img = Image.open(os.path.join(img_path,img_name))
+    try:        
+        img = Image.open(os.path.join(img_dir,img_name))
         img = transform_img(img)
         img = np.array(img)
         img = img[...,::-1]
@@ -60,22 +61,20 @@
                 mask = np.array(mask)
                 mosaic_area = impro.mask_area(mask)
             mosaic_img = mosaic.addmosaic_random(img, mask,'bounding') 
-            BOUNDING_flag = '_bound'
         else:
             mask = Image.open(os.path.join(ir_mask_path,random.choices(mask_names)[0]))
             mask = transform_mask(mask)
             mask = np.array(mask)
             mosaic_img = mosaic.addmosaic_random(img, mask)
-            BOUNDING_flag = ''    
-
-        if HD:#[128:384,128:384,:] --->256
-            cv2.imwrite(os.path.join(train_A_path,'%05d' % i+BOUNDING_flag+'.jpg'), mosaic_img)
-            cv2.imwrite(os.path.join(train_B_path,'%05d' % i+BOUNDING_flag+'.jpg'), img)
+
+        if MOD == 'HD':#[128:384,128:384,:] --->256
+            cv2.imwrite(os.path.join(train_A_path,'%05d' % i+suffix+'.jpg'), mosaic_img)
+            cv2.imwrite(os.path.join(train_B_path,'%05d' % i+suffix+'.jpg'), img)
         else:
             merge_img = impro.makedataset(mosaic_img, img)
-            cv2.imwrite(os.path.join(train_path,'%05d' % i+BOUNDING_flag+'.jpg'), merge_img)
+            cv2.imwrite(os.path.join(train_path,'%05d' % i+suffix+'.jpg'), merge_img)
         if MASK:
-            cv2.imwrite(os.path.join(mask_path,'%05d' % i+BOUNDING_flag+'.png'), mask)
-        print("Processing:",img_name," ","Remain:",len(img_names)-i)
+            cv2.imwrite(os.path.join(mask_path,'%05d' % i+suffix+'.png'), mask)
+        print('\r','Proc/all:'+str(i)+'/'+str(len(img_names)),util.get_bar(100*i/len(img_names),num=40),end='')
     except Exception as e:
         print(img_name,e)

models/loadmodel.py

@@ -2,23 +2,30 @@
 from .pix2pix_model import define_G
 from .pix2pixHD_model import define_G as define_G_HD
 from .unet_model import UNet
-from .video_model import HypoNet
+from .video_model import MosaicNet
+
+def show_paramsnumber(net,netname='net'):
+    parameters = sum(param.numel() for param in net.parameters())
+    parameters = round(parameters/1e6,2)
+    print(netname+' parameters: '+str(parameters)+'M')
+
 
 def pix2pix(opt):
     # print(opt.model_path,opt.netG)
     if opt.netG == 'HD':
         netG = define_G_HD(3, 3, 64, 'global' ,4)
     else:
         netG = define_G(3, 3, 64, opt.netG, norm='batch',use_dropout=True, init_type='normal', gpu_ids=[])
-
+    show_paramsnumber(netG,'netG')
     netG.load_state_dict(torch.load(opt.model_path))
     netG.eval()
     if opt.use_gpu:
         netG.cuda()
     return netG
 
 def video(opt):
-    netG = HypoNet(3*25+1, 3)
+    netG = MosaicNet(3*25+1, 3)
+    show_paramsnumber(netG,'netG')
     netG.load_state_dict(torch.load(opt.model_path))
     netG.eval()
     if opt.use_gpu:
@@ -28,6 +35,7 @@ def video(opt):
 
 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:
@@ -36,6 +44,7 @@ def unet_clean(opt):
 
 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:

models/video_model.py

@@ -151,9 +151,9 @@ def forward(self, x):
 
 
 
-class HypoNet(nn.Module):
+class MosaicNet(nn.Module):
     def __init__(self, in_channel, out_channel):
-        super(HypoNet, self).__init__()
+        super(MosaicNet, self).__init__()
 
         self.encoder_2d = encoder_2d(4,-1,64,n_blocks=9)
         self.encoder_3d = encoder_3d(in_channel)