Browse Source

Add -jpeg_chroma_subsampling_rate option for JPEG denoise training

nagadomi 9 years ago
parent
commit
d8ba661d6b
3 changed files with 12 additions and 16 deletions
  1. 9 13
      lib/pairwise_transform.lua
  2. 1 1
      lib/settings.lua
  3. 2 2
      train.lua

+ 9 - 13
lib/pairwise_transform.lua

@@ -122,10 +122,12 @@ function pairwise_transform.jpeg_(src, quality, size, offset, n, options)
    for i = 1, #quality do
       x = gm.Image(x, "RGB", "DHW")
       x:format("jpeg"):depth(8)
-      if options.jpeg_sampling_factors == 444 then
-	 x:samplingFactors({1.0, 1.0, 1.0})
-      else -- 420
+      if torch.uniform() < options.jpeg_chroma_subsampling_rate then
+	 -- YUV 420
 	 x:samplingFactors({2.0, 1.0, 1.0})
+      else
+	 -- YUV 444
+	 x:samplingFactors({1.0, 1.0, 1.0})
       end
       local blob, len = x:toBlob(quality[i])
       x:fromBlob(blob, len)
@@ -189,19 +191,13 @@ function pairwise_transform.jpeg(src, style, level, size, offset, n, options)
       end
    elseif style == "photo" then
       if level == 1 then
-	 return pairwise_transform.jpeg_(src, {torch.random(30, 75)},
+	 return pairwise_transform.jpeg_(src, {torch.random(70, 90)},
 					 size, offset, n,
 					 options)
       elseif level == 2 then
-	 if torch.uniform() > 0.6 then
-	    return pairwise_transform.jpeg_(src, {torch.random(30, 60)},
-					    size, offset, n, options)
-	 else
-	    local quality1 = torch.random(40, 60)
-	    local quality2 = quality1 - torch.random(5, 10)
-	    return pairwise_transform.jpeg_(src, {quality1, quality2},
-					    size, offset, n, options)
-	 end
+	 return pairwise_transform.jpeg_(src, {torch.random(50, 70)},
+					 size, offset, n,
+					 options)
       else
 	 error("unknown noise level: " .. level)
       end

+ 1 - 1
lib/settings.lua

@@ -37,7 +37,7 @@ cmd:option("-max_size", 256, 'if image is larger than max_size, image will be cr
 cmd:option("-batch_size", 8, 'mini batch size')
 cmd:option("-epoch", 200, 'number of total epochs to run')
 cmd:option("-thread", -1, 'number of CPU threads')
-cmd:option("-jpeg_sampling_factors", 444, '(444|420)')
+cmd:option("-jpeg_chroma_subsampling_rate", 0.0, 'the rate of YUV 4:2:0/YUV 4:4:4 in denoising training (0.0-1.0)')
 cmd:option("-validation_rate", 0.05, 'validation-set rate (number_of_training_images * validation_rate > 1)')
 cmd:option("-validation_crops", 80, 'number of cropping region per image in validation')
 cmd:option("-active_cropping_rate", 0.5, 'active cropping rate')

+ 2 - 2
train.lua

@@ -90,7 +90,7 @@ local function transformer(x, is_validation, n, offset)
    local active_cropping_rate = nil
    local active_cropping_tries = nil
    if is_validation then
-      active_cropping_rate = 0
+      active_cropping_rate = 0.0
       active_cropping_tries = 0
       random_color_noise_rate = 0.0
       random_overlay_rate = 0.0
@@ -126,7 +126,7 @@ local function transformer(x, is_validation, n, offset)
 					random_color_noise_rate = random_color_noise_rate,
 					random_overlay_rate = random_overlay_rate,
 					max_size = settings.max_size,
-					jpeg_sampling_factors = settings.jpeg_sampling_factors,
+					jpeg_chroma_subsampling_rate = settings.jpeg_chroma_subsampling_rate,
 					active_cropping_rate = active_cropping_rate,
 					active_cropping_tries = active_cropping_tries,
 					nr_rate = settings.nr_rate,