nagadomi · Dec 14, 2024 · Dec 14, 2024 · Dec 17, 2024 · Dec 17, 2024 · Dec 18, 2024
Showing with 77 additions and 9 deletions.

+42 −0 waifu2x/training/alpha_mask.py

+2 −2 waifu2x/training/create_training_data.py

+12 −2 waifu2x/training/dataset.py

+17 −1 waifu2x/training/trainer.py

+1 −1 waifu2x/ui_utils.py

+3 −3 waifu2x/utils.py
diff --git a/waifu2x/training/alpha_mask.py b/waifu2x/training/alpha_mask.py
@@ -0,0 +1,42 @@
+from PIL import Image, ImageEnhance
+
+def GenAlphaMask(image, grayscale, bg_color, darken):
+    alpha = image.getchannel("A")
+
+    if darken > 0:
+        enhancer = ImageEnhance.Brightness(image)
+        image = enhancer.enhance(float(darken))
+
+    background = Image.new("RGB", image.size, bg_color)
+    color = Image.composite(image.convert("RGB"), background, alpha)
+
+    if grayscale:
+        color = color.convert("L").convert("RGB")
+
+    return color, alpha.convert("RGB")
+
+
+def GenColorMask(image, grayscale, bg_color, darken):
+    alpha = image.getchannel("A")
+
+    if darken > 0:
+        enhancer = ImageEnhance.Brightness(image)
+        image = enhancer.enhance(float(darken))
+
+    flat_alpha = alpha.point(lambda p: 255 if p > 0 else 0)
+
+    background = Image.new("RGB", image.size, bg_color)
+    color = Image.composite(image.convert("RGB"), background, alpha)
+    mask = Image.composite(image.convert("RGB"), background, flat_alpha)
+
+    if grayscale:
+        color = color.convert("L").convert("RGB")
+        mask = mask.convert("L").convert("RGB")
+
+    return color, mask
+
+def GenMask(image, use_color, grayscale=False, bg_color=(255,255,255), darken=0.):
+    if use_color:
+        return GenColorMask(image, grayscale, bg_color, darken)
+    else:
+        return GenAlphaMask(image, grayscale, bg_color, darken)
diff --git a/waifu2x/training/create_training_data.py b/waifu2x/training/create_training_data.py
@@ -7,7 +7,7 @@
 import torchvision.transforms.functional as TF
 from torch.utils.data.dataset import Dataset
 from torch.utils.data import DataLoader
-from nunif.utils.pil_io import load_image_simple
+from nunif.utils.pil_io import load_image
 from nunif.utils.image_loader import list_images
 from nunif.transforms.std import pad
 
@@ -54,7 +54,7 @@ def __len__(self):
     def __getitem__(self, i):
         filename = self.files[i]
         bg_color = random.randint(0, 255)
-        im, _ = load_image_simple(filename, color="rgb", bg_color=bg_color)
+        im, _ = load_image(filename, color="rgb", keep_alpha=True)
         if im is None:
             return -1
         if self.args.pad:

diff --git a/waifu2x/training/dataset.py b/waifu2x/training/dataset.py
@@ -21,6 +21,7 @@
     shift_jpeg_block,
 )
 from .photo_noise import RandomPhotoNoiseX, add_validation_noise
+from .alpha_mask import GenMask
 from PIL.Image import Resampling
 
 
@@ -221,6 +222,10 @@ def __init__(self, input_dir,
                  model_offset,
                  scale_factor,
                  tile_size, num_samples=None,
+                 use_color=False,
+                 grayscale=False,
+                 bg_color=(255,255,255),
+                 darken=0.,
                  da_jpeg_p=0, da_scale_p=0, da_chshuf_p=0, da_unsharpmask_p=0,
                  da_grayscale_p=0, da_color_p=0, da_antialias_p=0, da_hflip_only=False,
                  da_no_rotate=False,
@@ -262,6 +267,10 @@ def __init__(self, input_dir,
         self.noise_level = noise_level
         self.model_offset = model_offset
         self.return_no_offset_y = return_no_offset_y
+        self.use_color = use_color
+        self.grayscale = grayscale
+        self.bg_color = bg_color
+        self.darken = darken
 
         if self.training:
             if noise_level >= 0:
@@ -380,7 +389,7 @@ def __init__(self, input_dir,
 
     def __getitem__(self, index):
         filename = super().__getitem__(index)
-        im, _ = pil_io.load_image_simple(filename, color="rgb")
+        im, _ = pil_io.load_image(filename, color="rgb", keep_alpha=True)
         if im is None:
             raise RuntimeError(f"Unable to load image: {filename}")
         if NEAREST_PREFIX in filename:
@@ -390,7 +399,8 @@ def __getitem__(self, index):
             x, y = self.transforms(im, im)
         else:
             im = self.gt_transforms(im)
-            x, y = self.transforms(im, im)
+            im_x, im_y = GenMask(im, use_color=self.use_color, grayscale=self.grayscale, bg_color=self.bg_color, darken=self.darken)
+            x, y = self.transforms(im_x, im_y)
 
         if not self.training:
             if self.noise_level >= 0:

diff --git a/waifu2x/training/trainer.py b/waifu2x/training/trainer.py
@@ -536,6 +536,10 @@ def create_dataloader(self, type):
                 noise_level=self.args.noise_level,
                 tile_size=self.args.size,
                 num_samples=self.args.num_samples,
+                use_color=self.args.use_color,
+                grayscale=self.args.grayscale,
+                bg_color=self.args.bg_color,
+                darken=self.args.darken,
                 da_jpeg_p=self.args.da_jpeg_p,
                 da_scale_p=self.args.da_scale_p,
                 da_chshuf_p=self.args.da_chshuf_p,
@@ -577,6 +581,10 @@ def create_dataloader(self, type):
                 style=self.args.style,
                 noise_level=self.args.noise_level,
                 tile_size=self.args.size,
+                use_color=self.args.use_color,
+                grayscale=self.args.grayscale,
+                bg_color=self.args.bg_color,
+                darken=self.args.darken,
                 fixed_deblur=self.args.fixed_deblur,
                 deblur=self.args.deblur,
                 resize_blur_range=self.args.resize_blur_range,
@@ -710,12 +718,20 @@ def register(subparsers, default_parser):
                         default="art",
                         help="image style used for jpeg noise level")
     parser.add_argument("--noise-level", type=int,
-                        choices=[0, 1, 2, 3],
+                        choices=[-1, 0, 1, 2, 3],
                         help="jpeg noise level for noise/noise_scale")
     parser.add_argument("--size", type=int, default=112,
                         help="input size")
     parser.add_argument("--num-samples", type=int, default=50000,
                         help="number of samples for each epoch")
+    parser.add_argument("--use-color", action="store_true",
+                        help="train on color bleed instead of alpha")
+    parser.add_argument("--grayscale", action="store_true",
+                        help="convert image to grayscale")
+    parser.add_argument("--bg-color", type=lambda s: tuple(map(int, s.split(','))), default=(255, 255, 255),
+                        help="background color as r,g,b (default: 255,255,255)")
+    parser.add_argument("--darken", type=float, default=0.0,
+                        help="darken input image")
     parser.add_argument("--drop-last", action="store_true",
                         help="force drop_last=True for DataLoader")
 

diff --git a/waifu2x/ui_utils.py b/waifu2x/ui_utils.py
@@ -185,7 +185,7 @@ def create_parser(required_true=True):
 
     parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument("--model-dir", type=str, help="model dir")
-    parser.add_argument("--noise-level", "-n", type=int, default=0, choices=[0, 1, 2, 3], help="noise level")
+    parser.add_argument("--noise-level", "-n", type=int, default=0, choices=[-1, 0, 1, 2, 3], help="noise level")
     parser.add_argument("--method", "-m", type=str,
                         choices=["scale4x", "scale2x",
                                  "noise_scale4x", "noise_scale2x",

diff --git a/waifu2x/utils.py b/waifu2x/utils.py
@@ -173,7 +173,7 @@ def _load_model(self, method, noise_level):
 
     def load_model(self, method, noise_level):
         assert (method in ("scale", "noise_scale", "noise", "scale4x", "noise_scale4x"))
-        assert (method in {"scale", "scale4x"} or 0 <= noise_level and noise_level < 4)
+        assert (method in {"scale", "scale4x"} or -1 <= noise_level and noise_level < 4)
 
         if method in {"scale", "scale4x", "noise"}:
             self._load_model(method, noise_level)
@@ -213,7 +213,7 @@ def load_model_all(self, load_4x=True):
 
     def render(self, x, method, noise_level, tile_size=None, batch_size=None, enable_amp=False):
         assert (method in ("scale", "noise_scale", "noise", "scale4x", "noise_scale4x"))
-        assert (method in {"scale", "scale4x"} or 0 <= noise_level and noise_level < 4)
+        assert (method in {"scale", "scale4x"} or -1 <= noise_level and noise_level < 4)
         if method == "scale":
             z = tiled_render(x, self.scale_model,
                              tile_size=tile_size, batch_size=batch_size,
@@ -255,7 +255,7 @@ def convert(self, x, alpha, method, noise_level,
         assert (x.shape[0] == 3)
         assert (alpha is None or alpha.shape[0] == 1 and alpha.shape[1:] == x.shape[1:])
         assert (method in ("scale", "scale4x", "noise_scale", "noise_scale4x", "noise"))
-        assert (method in {"scale", "scale4x"} or 0 <= noise_level and noise_level < 4)
+        assert (method in {"scale", "scale4x"} or -1 <= noise_level and noise_level < 4)
 
         if alpha is not None:
             # check all 1 alpha channel