added better support for template masks. letting you scale and rotate

Author: Blotz

data/steps_mask.png

@@ -34,7 +34,10 @@ def main() -> None:
         help="threshold for contrast. [-1.0, 1.0] Default: 1.0",
         default=1.0,
     )
-    threshold_template_group.add_argument("--template_path", type=str, help="path to template image")
+    template_group = threshold_template_group.add_argument_group("Template Options")
+    template_group.add_argument("--template_path", type=str, help="path to template image")
+    template_group.add_argument("--template_scale", type=float, help="scale up and down the template", default=1.0)
+    template_group.add_argument("--template_angle", type=float, help="angle of template", default=0)
 
     # optional arguments
     parser.add_argument(
@@ -79,14 +82,16 @@ def main() -> None:
     angle = args.angle
     threshold = args.threshold
     template_path = args.template_path
+    template_scale = args.template_scale
+    template_angle = args.template_angle
     sort_brightest = args.sort_brightest
     reversed_direction = args.reverse_sorting
     output_path = args.output
 
     # Manually check for mutual exclusive group
     if image_path is None and sys.stdin.isatty():
         parser.print_usage(sys.stderr)
-        parser.error("one of the arguments --image_path stdin is required")
+        parser.error("one of the arguments --image_path [stdin] is required")
 
         sys.exit(2)
 
@@ -116,7 +121,8 @@ def main() -> None:
             sys.exit(-1)
 
         template = cv2.imread(str(template_path))
-        create_mask = partial(image.create_template_mask, template)
+        template = image.scale_image(template, template_scale)
+        create_mask = partial(image.create_template_mask, template, template_angle-angle)
     else:
         create_mask = partial(image.create_contrast_mask, threshold)
 

src/pixelsort/__main__.py

@@ -126,7 +126,7 @@ def create_contrast_mask(threshold: float, image: np.ndarray) -> np.ndarray:
     return mask
 
 
-def create_template_mask(template: np.ndarray, image: np.ndarray) -> np.ndarray:
+def create_template_mask(template: np.ndarray, angle: float, image: np.ndarray) -> np.ndarray:
     """Create a mask for the given image using the template
 
     Args:
@@ -138,14 +138,17 @@ def create_template_mask(template: np.ndarray, image: np.ndarray) -> np.ndarray:
     """
     # remove bg color channels
     template = template[:, :, 0]
+    # Rotate Tile
+    template = scipy.ndimage.rotate(template, -angle, mode="constant")
+    
     # Tile the template to the size of the image
     horizontal_tiles = int(image.shape[1] / template.shape[1]) + 1
     vertical_tiles = int(image.shape[0] / template.shape[0]) + 1
     mask = np.tile(template, (vertical_tiles, horizontal_tiles))
 
     # Crop the mask to the size of the image
     mask = mask[: image.shape[0], : image.shape[1]]
-
+    
     return mask
 
 
@@ -182,5 +185,11 @@ def show_image(image: np.ndarray) -> None:
             break
     cv2.destroyAllWindows()
 
-def output_image(image: np.ndarray) -> None:
-    pass
+def scale_image(image: np.ndarray, scale: float) -> np.ndarray:
+    height, width, _ = image.shape
+    width *= scale
+    height *= scale
+    
+    dim = (int(width), int(height))
+    
+    return cv2.resize(image, dim, interpolation=cv2.INTER_NEAREST)