improved tiling

Author: HugoTini

__init__.py

@@ -79,7 +79,8 @@ class DeepBumpProperties(PropertyGroup):
         description='More overlap might help reducing some artifacts but takes longer to compute.',
         items=[('SMALL', 'Small', 'Small overlap between tiles.'),
                ('MEDIUM', 'Medium', 'Medium overlap between tiles.'),
-               ('BIG', 'Big', 'Big overlap between tiles.')]
+               ('LARGE', 'Large', 'Large overlap between tiles.')],
+        default='LARGE'
     )
 
 
@@ -136,13 +137,13 @@ def execute(self, context):
         img = np.mean(img[0:3], axis=0, keepdims=True)
 
         # Split image in tiles
-        tile_size = (256, 256)
-        OVERLAP = context.scene.deep_bump_tool.tiles_overlap_enum
-        overlaps = {'SMALL': 20, 'MEDIUM': 50, 'BIG': 124}
-        stride_size = (tile_size[0]-overlaps[OVERLAP],
-                       tile_size[1]-overlaps[OVERLAP])
         print('DeepBump : tilling')
-        tiles, paddings = infer.tiles_split(img, tile_size, stride_size)
+        tile_size = 256
+        OVERLAP = context.scene.deep_bump_tool.tiles_overlap_enum
+        overlaps = {'SMALL': tile_size//6, 'MEDIUM': tile_size//4, 'LARGE': tile_size//2}
+        stride_size = tile_size-overlaps[OVERLAP]
+        tiles, paddings = infer.tiles_split(img, (tile_size, tile_size),
+                                            (stride_size, stride_size))
 
         # Load model (if not already loaded)
         if self.ort_session is None:
@@ -160,8 +161,8 @@ def execute(self, context):
 
         # Merge tiles
         print('DeepBump : merging')
-        pred_img = infer.tiles_merge(
-            pred_tiles, stride_size, (3, img.shape[1], img.shape[2]), paddings)
+        pred_img = infer.tiles_merge(pred_tiles, (stride_size, stride_size), 
+                                    (3, img.shape[1], img.shape[2]), paddings)
 
         # Normalize each pixel to unit vector
         pred_img = infer.normalize(pred_img)

cli.md

@@ -21,4 +21,4 @@ Go to the DeepBump folder, then use with :
 
     python deepbump_cli.py -i <path_input_image> -n <path_generated_normal> -o <overlap>
 
-where `<overlap>` can be `small`, `medium` or `big`.
+where `<overlap>` can be `small`, `medium` or `large`.

deepbump_cli.py

@@ -1,51 +1,57 @@
-from PIL import Image
-import numpy as np
-import infer
-import onnxruntime as ort
-from argparse import ArgumentParser
-
-# Disable MS telemetry
-ort.disable_telemetry_events()
-
-# Parse cli arguments
-parser = ArgumentParser()
-parser.add_argument('-i', '--img_path', dest='img_path', required=True,
-                    help='path to the input image')
-parser.add_argument('-n', '--normal_path', dest='normal_path', required=True,
-                    help='path where the generated normal map is saved')
-parser.add_argument('-o', '--overlap', dest='overlap', required=True,
-                    help='tiles overlap size, must be "small", "medium" or "big"',
-                    choices=['small', 'medium', 'big'])
-args = parser.parse_args()
-
-# Load image
-img = np.array(Image.open(args.img_path)) / 255.0
-# Reshape to C,H,W
-img = np.transpose(img, (2, 0, 1))
-# Grayscale
-img = np.mean(img[0:3], axis=0, keepdims=True)
-
-# Split image in tiles
-print('tilling')
-tile_size = (256, 256)
-overlaps = {'small': 20, 'medium': 50, 'big': 124}
-stride_size = (tile_size[0]-overlaps[args.overlap], tile_size[1]-overlaps[args.overlap])
-tiles, paddings = infer.tiles_split(img, tile_size, stride_size)
-
-# Predict tiles normal map
-print('generating')
-
-def progress_print(current, total):
-    print('{}/{}'.format(current, total))
-    
-ort_session = ort.InferenceSession('./deepbump256.onnx')
-pred_tiles = infer.tiles_infer(tiles, ort_session, progress_callback=progress_print)
-
-# Merge tiles
-print('merging')
-pred_img = infer.tiles_merge(pred_tiles, stride_size, (3, img.shape[1], img.shape[2]), paddings)
-
-# Save image 
-pred_img = pred_img.transpose((1,2,0))
-pred_img = Image.fromarray((pred_img*255.0).astype(np.uint8))
-pred_img.save(args.normal_path)  
+from PIL import Image
+import numpy as np
+import infer
+import onnxruntime as ort
+from argparse import ArgumentParser
+
+
+def progress_print(current, total):
+    print('{}/{}'.format(current, total))
+
+
+# Disable MS telemetry
+ort.disable_telemetry_events()
+
+# Parse cli arguments
+parser = ArgumentParser()
+parser.add_argument('-i', '--img_path', dest='img_path', required=True,
+                    help='path to the input image')
+parser.add_argument('-n', '--normal_path', dest='normal_path', required=True,
+                    help='path where the generated normal map is saved')
+parser.add_argument('-o', '--overlap', dest='overlap', required=True,
+                    help='tiles overlap size, must be "small", "medium" or "large"',
+                    choices=['small', 'medium', 'large'])
+args = parser.parse_args()
+
+# Load image
+img = np.array(Image.open(args.img_path)) / 255.0
+# Reshape to C,H,W
+img = np.transpose(img, (2, 0, 1))
+# Grayscale
+img = np.mean(img[0:3], axis=0, keepdims=True)
+
+# Split image in tiles
+print('tilling')
+tile_size = 256
+overlaps = {'small':  tile_size//6, 'medium':  tile_size//4,
+            'large': tile_size//2}
+stride_size = tile_size-overlaps[args.overlap]
+tiles, paddings = infer.tiles_split(img, (tile_size, tile_size),
+                                    (stride_size, stride_size))
+
+# Predict tiles normal map
+print('generating')
+ort_session = ort.InferenceSession('./deepbump256.onnx')
+pred_tiles = infer.tiles_infer(tiles, ort_session,
+                               progress_callback=progress_print)
+
+# Merge tiles
+print('merging')
+pred_img = infer.tiles_merge(pred_tiles, (stride_size, stride_size),
+                             (3, img.shape[1], img.shape[2]),
+                             paddings)
+
+# Save image
+pred_img = pred_img.transpose((1, 2, 0))
+pred_img = Image.fromarray((pred_img*255.0).astype(np.uint8))
+pred_img.save(args.normal_path)

infer.py

@@ -14,11 +14,11 @@ def tiles_split(img, tile_size, stride_size):
     img_h, img_w = img.shape[1], img.shape[2]
 
     # stride must be even
-    assert((stride_h % 2 == 0) and (stride_w % 2 == 0))
-    # stride must greater than half tile
-    assert((stride_h > tile_h/2) and (stride_w > tile_w/2))
-    # stride must be smaller than tile size
-    assert((stride_h < tile_h) and (stride_w < tile_w))
+    assert (stride_h % 2 == 0) and (stride_w % 2 == 0)
+    # stride must be greater or equal than half tile
+    assert (stride_h >= tile_h/2) and (stride_w >= tile_w/2)
+    # stride must be smaller or equal tile size
+    assert (stride_h <= tile_h) and (stride_w <= tile_w)
 
     # find total height & width padding sizes
     pad_h, pad_w = 0, 0
@@ -97,11 +97,10 @@ def generate_mask(tile_size, stride_size):
     mask[-ramp_h:, ramp_w:-ramp_w] = np.transpose(np.linspace(1, 0, num=ramp_h)[None],
                                                   (1, 0))
 
+    # Assume tiles are squared
+    assert ramp_h == ramp_w
     # top left corner
-    plane_h = np.repeat(np.transpose(np.linspace(0, 0.25, num=ramp_h)[None], (1, 0)),
-                        ramp_w, axis=1)
-    plane_w = np.repeat(np.linspace(0, 0.25, num=ramp_w)[None], ramp_h, axis=0)
-    corner = plane_h + plane_w
+    corner = np.rot90(corner_mask(ramp_h), 2)
     mask[:ramp_h, :ramp_w] = corner
     # top right corner
     corner = np.flip(corner, 1)
@@ -116,6 +115,44 @@ def generate_mask(tile_size, stride_size):
     return mask
 
 
+def corner_mask(side_length):
+    '''Generates the corner part of the pyramidal-like mask. 
+    Currently, only for square shapes.'''
+
+    corner = np.zeros([side_length, side_length])
+
+    for h in range(0, side_length):
+        for w in range(0, side_length):
+            if h >= w:
+                sh = h / (side_length-1)
+                corner[h, w] = 1-sh
+            if h <= w:
+                sw = w / (side_length-1)
+                corner[h, w] = 1-sw
+
+    return corner-0.25*scaling_mask(side_length)
+
+
+def scaling_mask(side_length):
+
+    scaling = np.zeros([side_length, side_length])
+
+    for h in range(0, side_length):
+        for w in range(0, side_length):
+            sh = h / (side_length-1)
+            sw = w / (side_length-1)
+            if h >= w and h <= side_length-w:
+                scaling[h, w] = sw
+            if h <= w and h <= side_length-w:
+                scaling[h, w] = sh
+            if h >= w and h >= side_length-w:
+                scaling[h, w] = 1-sh
+            if h <= w and h >= side_length-w:
+                scaling[h, w] = 1-sw
+
+    return 2*scaling
+
+
 def tiles_merge(tiles, stride_size, img_size, paddings):
     '''Merges the list of tiles into one image. img_size is the original size, before 
     padding.'''
@@ -127,11 +164,11 @@ def tiles_merge(tiles, stride_size, img_size, paddings):
     stride_h, stride_w = stride_size
 
     # stride must be even
-    assert((stride_h % 2 == 0) and (stride_w % 2 == 0))
-    # stride must greater than half tile
-    assert((stride_h > tile_h/2) and (stride_w > tile_w/2))
-    # stride must be smaller than tile size
-    assert((stride_h < tile_h) and (stride_w < tile_w))
+    assert (stride_h % 2 == 0) and (stride_w % 2 == 0)
+    # stride must be greater or equal than half tile
+    assert (stride_h >= tile_h/2) and (stride_w >= tile_w/2)
+    # stride must be smaller or equal tile size
+    assert (stride_h <= tile_h) and (stride_w <= tile_w)
 
     merged = np.zeros((img_size[0], height, width))
     mask = generate_mask((tile_h, tile_w), stride_size)