Fix bg-channel bug in s_self_tile and s_self_tile_majority (unlocks Tasks 1, 304)

own-solver/neurogolf_solver/solvers/wave3.py

@@ -14,7 +14,7 @@ from ..constants import GH, GW
 
 
 def s_self_tile(td):
-    """Self-tiling: each non-bg pixel → full copy of input. Output = IH^2 x IW^2."""
+    """Self-tiling: each non-bg pixel -> full copy of input. Output = IH^2 x IW^2."""
     exs = get_exs(td)
     sp = fixed_shapes(td)
     if sp is None:
@@ -46,6 +46,9 @@ def s_self_tile(td):
         _make_int64_init('tile_rp', [1, 1, IH, IW]),
         numpy_helper.from_array(np.array([0,0,0,0,0,0,pad_h,pad_w], dtype=np.int64), 'pads'),
         numpy_helper.from_array(np.array(0.0, dtype=np.float32), 'pad_cv'),
+        numpy_helper.from_array(np.float32(0.5), 'half'),
+        numpy_helper.from_array(np.ones((1, 1, OH, OW), dtype=np.float32), 'grid_mask'),
+        numpy_helper.from_array(np.zeros((1, 9, OH, OW), dtype=np.float32), 'z9'),
     ]
     nodes = [
         helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cropped']),
@@ -54,7 +57,14 @@ def s_self_tile(td):
         helper.make_node('ConvTranspose', ['mask_small', 'up_k'], ['mask_big'],
                         kernel_shape=[IH, IW], strides=[IH, IW]),
         helper.make_node('Tile', ['cropped', 'tile_rp'], ['tiled']),
-        helper.make_node('Mul', ['tiled', 'mask_big'], ['result']),
+        helper.make_node('Mul', ['tiled', 'mask_big'], ['result_raw']),
+        # bg-channel fix: set ch0=1 for bg positions within OH x OW grid
+        helper.make_node('ReduceMax', ['result_raw'], ['has_color'], axes=[1], keepdims=1),
+        helper.make_node('Greater', ['has_color', 'half'], ['has_color_b']),
+        helper.make_node('Cast', ['has_color_b'], ['has_color_f'], to=TensorProto.FLOAT),
+        helper.make_node('Sub', ['grid_mask', 'has_color_f'], ['bg_mask']),
+        helper.make_node('Concat', ['bg_mask', 'z9'], ['bg_oh'], axis=1),
+        helper.make_node('Add', ['result_raw', 'bg_oh'], ['result']),
         helper.make_node('Pad', ['result', 'pads', 'pad_cv'], ['output'], mode='constant'),
     ]
     return mk(nodes, inits)
@@ -94,6 +104,9 @@ def s_self_tile_majority(td):
         _make_int64_init('tile_rp', [1, 1, IH, IW]),
         numpy_helper.from_array(np.array([0,0,0,0,0,0,pad_h,pad_w], dtype=np.int64), 'pads'),
         numpy_helper.from_array(np.array(0.0, dtype=np.float32), 'pad_cv'),
+        numpy_helper.from_array(np.float32(0.5), 'half'),
+        numpy_helper.from_array(np.ones((1, 1, OH, OW), dtype=np.float32), 'grid_mask'),
+        numpy_helper.from_array(np.zeros((1, 9, OH, OW), dtype=np.float32), 'z9'),
     ]
     nodes = [
         helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cropped']),
@@ -106,7 +119,14 @@ def s_self_tile_majority(td):
         helper.make_node('ConvTranspose', ['mask_small', 'up_k'], ['mask_big'],
                         kernel_shape=[IH, IW], strides=[IH, IW]),
         helper.make_node('Tile', ['cropped', 'tile_rp'], ['tiled']),
-        helper.make_node('Mul', ['tiled', 'mask_big'], ['result']),
+        helper.make_node('Mul', ['tiled', 'mask_big'], ['result_raw']),
+        # bg-channel fix: set ch0=1 for bg positions within OH x OW grid
+        helper.make_node('ReduceMax', ['result_raw'], ['has_color'], axes=[1], keepdims=1),
+        helper.make_node('Greater', ['has_color', 'half'], ['has_color_b']),
+        helper.make_node('Cast', ['has_color_b'], ['has_color_f'], to=TensorProto.FLOAT),
+        helper.make_node('Sub', ['grid_mask', 'has_color_f'], ['bg_mask']),
+        helper.make_node('Concat', ['bg_mask', 'z9'], ['bg_oh'], axis=1),
+        helper.make_node('Add', ['result_raw', 'bg_oh'], ['result']),
         helper.make_node('Pad', ['result', 'pads', 'pad_cv'], ['output'], mode='constant'),
     ]
     return mk(nodes, inits)