Update to Pytorch 1.0

CleanMosaic.py

@@ -17,7 +17,7 @@ opt = CleanOptions().getparse()
 def get_mosaic_position(img_origin):
     mask =runmodel.run_unet_rectim(img_origin,net_mosaic_pos,use_gpu = opt.use_gpu)
     mask = impro.mask_threshold(mask,10,128)
-    x,y,size,area = impro.boundingSquare(mask,threshold=128,Ex_mul=1.5)
+    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

README.md

@@ -10,9 +10,9 @@ The code do not include the part of training, I will finish it in my free time.
 
 ## Prerequisites
 - Linux, (I didn't try this code on Windows or Mac OS)
-- Python 3.5+
+- Python 3.6+
 - ffmpeg
-- Pytorch 0.4   [(Pytorch 1.0+ is available)](https://github.com/HypoX64/DeepMosaics)
+- Pytorch 1.0+  [(Old version codes)](https://github.com/HypoX64/DeepMosaics/tree/Pytorch0.4)
 - CPU or NVIDIA GPU + CUDA CuDNN
 
 ## Getting Started
@@ -23,8 +23,8 @@ 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=1PXt3dE9Eez2xUqpemLJutwTCC0tW-D2g)
- [[百度云,提取码z8vz]](https://pan.baidu.com/s/1Wi8T6PE4ExTjrHVQhv3rJA) 
+[[Google Drive]](https://drive.google.com/open?id=10nARsiZoZGcaKw40nQu9fJuRp1oeabPs)
+ [[百度云,提取码7thu]](https://pan.baidu.com/s/1IG4bdIiIC9PH9-oEyae5Sg) 
 
 ### Dependencies
 This code depends on numpy, scipy, opencv-python, torchvision, available via pip install.

models/loadmodel.py

@@ -4,9 +4,7 @@ from .unet_model import UNet
 
 def pix2pix(model_path,G_model_type,use_gpu = True):
     gpu_ids=[]
-    if use_gpu:
-        gpu_ids=[0]
-    netG = define_G(3, 3, 64, G_model_type, norm='instance', init_type='normal', gpu_ids=gpu_ids)
+    netG = define_G(3, 3, 64, G_model_type, norm='batch', init_type='normal', gpu_ids=gpu_ids)
     netG.load_state_dict(torch.load(model_path))
     netG.eval()
     if use_gpu:
@@ -20,6 +18,3 @@ def unet(model_path,use_gpu = True):
     if use_gpu:
         net.cuda()
     return net
-
-
-# def unet():
\ No newline at end of file

models/unet_model.py

@@ -29,4 +29,4 @@ class UNet(nn.Module):
         x = self.up3(x, x2)
         x = self.up4(x, x1)
         x = self.outc(x)
-        return F.sigmoid(x)
+        return torch.sigmoid(x)

models/unet_parts.py

@@ -45,7 +45,14 @@ class down(nn.Module):
         x = self.mpconv(x)
         return x
 
+class Upsample(nn.Module):
+    def __init__(self,  scale_factor):
+        super(Upsample, self).__init__()
+        self.scale_factor = scale_factor
+    def forward(self, x):
+        return F.interpolate(x, scale_factor=self.scale_factor,mode='bilinear', align_corners=True)
 
+F.interpolate
 class up(nn.Module):
     def __init__(self, in_ch, out_ch, bilinear=True):
         super(up, self).__init__()
@@ -53,7 +60,7 @@ class up(nn.Module):
         #  would be a nice idea if the upsampling could be learned too,
         #  but my machine do not have enough memory to handle all those weights
         if bilinear:
-            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.up = Upsample(scale_factor=2)
         else:
             self.up = nn.ConvTranspose2d(in_ch//2, in_ch//2, 2, stride=2)
 

options/addmosaic_options.py

@@ -18,7 +18,7 @@ class AddOptions():
         self.parser.add_argument('--mask_extend', type=int, default=20,help='more mosaic area')
         self.parser.add_argument('--mask_threshold', type=int, default=64,help='threshold of recognize mosaic position 0~255')
         self.parser.add_argument('--output_size', type=int, default=0,help='size of output file,if 0 -> origin')
-        self.parser.add_argument('--tempimage_type', type=str, default='jpg',help='type of temp image, png | jpg, png is better but occupy more storage space')
+        self.parser.add_argument('--tempimage_type', type=str, default='png',help='type of temp image, png | jpg, png is better but occupy more storage space')
         self.initialized = True
 
     def getparse(self):

options/cleanmosaic_options.py

@@ -18,7 +18,7 @@ class CleanOptions():
                                 help='name of model use to find mosaic position')
         self.parser.add_argument('--no_feather', action='store_true', help='if true, no edge feather,but run faster')
         self.parser.add_argument('--medfilt_num', type=int, default=11,help='medfilt window of mosaic movement in the video')
-        self.parser.add_argument('--tempimage_type', type=str, default='jpg',help='type of temp image, png | jpg, png is better but occupy more storage space')
+        self.parser.add_argument('--tempimage_type', type=str, default='png',help='type of temp image, png | jpg, png is better but occupy more storage space')
 #       self.parser.add_argument('--zoom_multiple', type=float, default=1.0,help='zoom video')
         self.initialized = True
 

util/image_processing.py

@@ -4,7 +4,7 @@ import numpy as np
 
 def resize(img,size):
     h, w = img.shape[:2]
-    if min(h, w) ==size:
+    if np.min((w,h)) ==size:
         return img
     if w >= h:
         res = cv2.resize(img,(int(size*w/h), size))
@@ -62,7 +62,7 @@ def mergeimage(img1,img2,orgin_image):
     result_img = cv2.add(new_img1,new_img2)
     return result_img
 
-def boundingSquare(mask,threshold,Ex_mul):
+def boundingSquare(mask,Ex_mul):
     # thresh = mask_threshold(mask,10,threshold)
     area = mask_area(mask)
     if area == 0 :

util/mosaic.py

@@ -6,6 +6,7 @@ from .image_processing import resize,ch_one2three,mask_area
 
 
 def addmosaic(img,mask,n,out_size = 0,model = 'squa_avg'):
+    n = int(n)
     if out_size:
         img = resize(img,out_size)      
     h, w = img.shape[:2]
@@ -59,22 +60,24 @@ def random_mosaic_mod(img,mask,n):
     return img
 
 def random_mosaic(img,mask):
-    img = resize(img,512)
+    # img = resize(img,512)
     h,w = img.shape[:2]
     mask = cv2.resize(mask,(w,h))
+    alpha = np.min((w,h))/512
     #area_avg=5925*4
     try:
         area = mask_area(mask)
     except:
         area = 0
+    area = area/(alpha*alpha)
     if area>50000:
-        img_mosaic = random_mosaic_mod(img,mask,random.randint(14,26))
+        img_mosaic = random_mosaic_mod(img,mask,alpha*random.uniform(16,28))
     elif 20000<area<=50000:
-        img_mosaic = random_mosaic_mod(img,mask,random.randint(10,18))
+        img_mosaic = random_mosaic_mod(img,mask,alpha*random.uniform(12,20))
     elif 5000<area<=20000:
-        img_mosaic = random_mosaic_mod(img,mask,random.randint(8,14))
+        img_mosaic = random_mosaic_mod(img,mask,alpha*random.uniform(8,15))
     elif 0<=area<=5000:
-        img_mosaic = random_mosaic_mod(img,mask,random.randint(4,8))
+        img_mosaic = random_mosaic_mod(img,mask,alpha*random.uniform(4,10))
     else:
         pass
     return img_mosaic