Merge pull request #24 from loiccoyle/feat/gpu

feat: use gpu to compute the distance matrix

README.md

@@ -20,6 +20,10 @@ In a terminal:
 
 ```sh
 pip install phomo
+
+# or for GPU acceleration:
+
+pip install 'phomo[cuda]'
 ```
 
 As always, it is usually a good idea to use a [virtual environment](https://docs.python.org/3/library/venv.html).
@@ -39,7 +43,7 @@ Once it is installed, you can use the `phomo` command.
 It would go something like:
 
 ```sh
-phomo master.png tile_directory -S 20 20 -o mosaic.png
+phomo master.png tile_directory/ -S 20 20 -o mosaic.png
 ```
 
 If in doubt see the help:
@@ -50,9 +54,9 @@ If in doubt see the help:
 $ phomo -h
 usage: phomo [-h] [-o OUTPUT] [-c MASTER_CROP_RATIO]
              [-s MASTER_SIZE [MASTER_SIZE ...]] [-C TILE_CROP_RATIO]
-             [-S TILE_SIZE [TILE_SIZE ...]] [-n N_APPEARANCES] [-v] [-b] [-g]
-             [-d SUBDIVISIONS [SUBDIVISIONS ...]]
-             [-m {greyscale,norm,luv_approx}] [-j WORKERS]
+             [-S TILE_SIZE [TILE_SIZE ...]] [-n N_APPEARANCES] [-b] [-g]
+             [-d SUBDIVISIONS [SUBDIVISIONS ...]] [-G]
+             [-m {greyscale,norm,luv_approx}] [-j WORKERS] [-v]
              master tile_dir
 
 positional arguments:
@@ -73,17 +77,19 @@ options:
                         Resize tile images to width, height.
   -n N_APPEARANCES, --n-appearances N_APPEARANCES
                         The number of times a tile can appear in the mosaic.
-  -v, --verbose         Verbosity.
   -b, --black-and-white
                         Convert master and tile images to black and white.
   -g, --show-grid       Show the tile grid, don't build the mosiac.
   -d SUBDIVISIONS [SUBDIVISIONS ...], --subdivisions SUBDIVISIONS [SUBDIVISIONS ...]
                         Grid subdivision thresholds.
+  -G, --gpu             Use GPU for distance matrix computation. Requires
+                        installing with `pip install 'phomo[cuda]'`.
   -m {greyscale,norm,luv_approx}, --metric {greyscale,norm,luv_approx}
                         Distance metric.
   -j WORKERS, --workers WORKERS
                         Number of workers use to run when computing the
                         distance matrix.
+  -v, --verbose         Verbosity.
 ```
 
 <!-- help end -->

phomo/__main__.py

@@ -51,7 +51,6 @@ def parse_args(args: List[str]) -> argparse.Namespace:
         nargs="+",
         default=None,
     )
-
     parser.add_argument(
         "-C",
         "--tile-crop-ratio",
@@ -74,13 +73,6 @@ def parse_args(args: List[str]) -> argparse.Namespace:
         type=int,
         default=1,
     )
-    parser.add_argument(
-        "-v",
-        "--verbose",
-        help="Verbosity.",
-        action="count",
-        default=0,
-    ),
     parser.add_argument(
         "-b",
         "--black-and-white",
@@ -101,6 +93,12 @@ def parse_args(args: List[str]) -> argparse.Namespace:
         default=[],
         type=float,
     )
+    parser.add_argument(
+        "-G",
+        "--gpu",
+        help="Use GPU for distance matrix computation. Requires installing with `pip install 'phomo[cuda]'`.",
+        action="store_true",
+    )
     parser.add_argument(
         "-m",
         "--metric",
@@ -116,6 +114,13 @@ def parse_args(args: List[str]) -> argparse.Namespace:
         default=1,
         type=int,
     )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        help="Verbosity.",
+        action="count",
+        default=0,
+    )
     return parser.parse_args(args)
 
 
@@ -180,7 +185,12 @@ def main():
         grid_im = mosaic.grid.plot()
         grid_im.show()
     else:
-        mosaic_im = mosaic.build(workers=args.workers, metric=args.metric)
+        d_matrix = (
+            mosaic.d_matrix(metric=args.metric, workers=args.workers)
+            if not args.gpu
+            else mosaic.d_matrix_cuda(metric=args.metric)
+        )
+        mosaic_im = mosaic.build(d_matrix)
         if args.output is None:
             mosaic_im.show()
         else:

phomo/metrics.py

@@ -6,8 +6,9 @@
 
 
 class MetricCallable(Protocol):  # type: ignore
-    def __call__(self, master_chunk: np.ndarray, tile_arrays: np.ndarray) -> np.ndarray:
-        ...
+    def __call__(
+        self, master_chunk: np.ndarray, tile_arrays: np.ndarray
+    ) -> np.ndarray: ...
 
 
 def register_metric(func):
@@ -107,6 +108,7 @@ def luv_approx(
             arrays.
     """
     r = (master_chunk[:, :, 0] + tile_arrays[:, :, :, 0]) // 2
+    r = r.astype(float)
     d = np.subtract(master_chunk, tile_arrays, dtype=float)
     return np.linalg.norm(
         (

phomo/mosaic.py

@@ -1,12 +1,13 @@
 import logging
+import math
 from functools import partial
 from multiprocessing.pool import Pool as MpPool
-from typing import Optional, Tuple, Union
+from typing import Tuple, Union
 
 import numpy as np
 from PIL import Image
-from tqdm.auto import tqdm
 from scipy.optimize import linear_sum_assignment
+from tqdm.auto import tqdm
 
 from .grid import Grid
 from .master import Master
@@ -37,7 +38,8 @@ def __init__(
 
             >>> pool = Pool.from_dir("tiles")
             >>> master = Master.from_file("master.png")
-            >>> Mosaic(master, pool, n_appearances=1).build()
+            >>> mosaic = Mosaic(master, pool, n_appearances=1)
+            >>> mosaic.build(mosaic.d_matrix())
         """
         self._log = logging.getLogger(__name__)
         self.master = master
@@ -82,7 +84,7 @@ def _d_matrix_worker(
             array = resize_array(array, (self.tile_shape[1], self.tile_shape[0]))
         return metric_func(array, self.pool.arrays, **kwargs)
 
-    def compute_d_matrix(
+    def d_matrix(
         self,
         workers: int = 1,
         metric: Union[str, MetricCallable] = "norm",
@@ -144,13 +146,79 @@ def compute_d_matrix(
         self._log.debug("d_matrix shape: %s", d_matrix.shape)
         return d_matrix
 
-    def build_greedy(
-        self,
-        workers: int = 1,
-        metric: Union[str, MetricCallable] = "norm",
-        d_matrix: Optional[np.ndarray] = None,
-        **kwargs,
-    ) -> Image.Image:
+    def d_matrix_cuda(self, metric: str = "norm") -> np.ndarray:
+        """Compute the distance matrix using CUDA for GPU acceleration.
+
+        Args:
+            metric: The distance metric used for the distance matrix. Either "norm" or "greyscale".
+
+        Returns:
+            Distance matrix, shape: (number of master arrays, number of tiles in the pool).
+        """
+
+        try:
+            from numba import cuda
+        except ImportError:
+            raise ImportError(
+                "Numba is required for CUDA support, run \"pip install 'phomo[cuda]'\" to install it."
+            )
+
+        if metric not in ["norm", "greyscale"]:
+            raise ValueError(
+                f"Invalid metric '{metric}'. When using gpu `metric' must be 'norm' or 'greyscale'."
+            )
+
+        self._log.info("Computing distance matrix with CUDA.")
+
+        # when the grid has been subdivided the master arrays will be smaller, so we grow them to match
+        # the tile size
+        grid_arrays = [
+            array
+            if array.shape == self.tile_shape
+            else resize_array(array, self.tile_shape)
+            for array in self.grid.arrays
+        ]
+        pool_arrays = self.pool.arrays
+        if metric == "greyscale":
+            grid_arrays = [array.sum(axis=-1, keepdims=True) for array in grid_arrays]
+            pool_arrays = [array.sum(axis=-1, keepdims=True) for array in pool_arrays]
+
+        # Transfer the master and pool arrays to the GPU.
+        master_arrays_device = cuda.to_device(grid_arrays)
+        pool_arrays_device = cuda.to_device(pool_arrays)
+
+        # Allocate memory for the distance matrix on the GPU.
+        d_matrix_device = cuda.device_array((len(grid_arrays), len(pool_arrays)))
+
+        # Define the CUDA kernel for computing the distance matrix.
+        @cuda.jit
+        def compute_d_matrix_kernel(master_arrays, pool_arrays, d_matrix):
+            i, j = cuda.grid(2)  # type: ignore
+            if i < master_arrays.shape[0] and j < pool_arrays.shape[0]:
+                distance = 0.0
+                for x in range(master_arrays.shape[1]):
+                    for y in range(master_arrays.shape[2]):
+                        for c in range(master_arrays.shape[3]):
+                            diff = master_arrays[i, x, y, c] - pool_arrays[j, x, y, c]
+                            distance += diff * diff
+                d_matrix[i, j] = math.sqrt(distance)
+
+        # Define the number of threads per block and blocks per grid.
+        threads_per_block = (16, 16)
+        blocks_per_grid_x = math.ceil(len(grid_arrays) / threads_per_block[0])
+        blocks_per_grid_y = math.ceil(len(pool_arrays) / threads_per_block[1])
+        blocks_per_grid = (blocks_per_grid_x, blocks_per_grid_y)
+
+        # Launch the kernel.
+        compute_d_matrix_kernel[blocks_per_grid, threads_per_block](  # type: ignore
+            master_arrays_device, pool_arrays_device, d_matrix_device
+        )
+
+        self._log.debug("d_matrix shape: %s", d_matrix_device.shape)
+        # Copy the result back to the host.
+        return d_matrix_device.copy_to_host()
+
+    def build_greedy(self, d_matrix: np.ndarray) -> Image.Image:
         """Construct the mosaic image using a greedy tile assignement algorithm.
 
         Args:
@@ -167,10 +235,6 @@ def build_greedy(
         """
         mosaic = np.zeros((self.size[1], self.size[0], 3))
 
-        # Compute the distance matrix.
-        if d_matrix is None:
-            d_matrix = self.compute_d_matrix(workers=workers, metric=metric, **kwargs)
-
         # Keep track of tiles and sub arrays.
         placed_master_arrays = set()
         placed_tiles = set()
@@ -207,34 +271,27 @@ def build_greedy(
         pbar.close()
         return Image.fromarray(np.uint8(mosaic))
 
-    def build(
-        self,
-        workers: int = 1,
-        metric: Union[str, MetricCallable] = "norm",
-        d_matrix: Optional[np.ndarray] = None,
-        **kwargs,
-    ) -> Image.Image:
+    def build(self, d_matrix: np.ndarray) -> Image.Image:
         """Construct the mosaic image by solving the linear sum assignment problem.
         See: https://en.wikipedia.org/wiki/Assignment_problem
 
         Args:
-            workers: The number of workers to use when computing the
-                distance matrix.
-            metric: The distance metric used for the distance matrix. Either
-                provide a string, for implemented metrics see ``phomo.metrics.METRICS``.
-                Or a callable, which should take two ``np.ndarray``s and return a float.
-            d_matrix: Use a pre-computed distance matrix.
-            **kwargs: Passed to the `metric` function.
+            d_matrix: The distanace matrix to use to build the mosaic.
 
         Returns:
             The PIL.Image instance of the mosaic.
+
+        Examples:
+            Building a mosaic.
+
+            >>> mosaic.build(mosaic.d_matrix())
+
+            On a GPU.
+
+            >>> mosaic.build(mosaic.d_matrix_cuda())
         """
         mosaic = np.zeros((self.size[1], self.size[0], 3))
 
-        # Compute the distance matrix.
-        if d_matrix is None:
-            d_matrix = self.compute_d_matrix(workers=workers, metric=metric, **kwargs)
-
         # expand the dmatrix to allow for repeated tiles
         if self.n_appearances > 0:
             d_matrix = np.tile(d_matrix, self.n_appearances)

phomo/palette.py

@@ -8,7 +8,8 @@
 class Palette:
     """Colour palette methods."""
 
-    pixels: np.ndarray
+    @property
+    def pixels(self) -> np.ndarray: ...
 
     def palette(self, **kwargs) -> Tuple[np.ndarray, np.ndarray]:
         """Compute the colour distribution.

phomo/utils.py

@@ -111,7 +111,8 @@ def open_img_file(
         Image instance.
     """
     with Image.open(img_file) as img:
-        img = exif_transpose(img)
+        img_t = exif_transpose(img)
+        img = img_t if img_t is not None else img
         if crop_ratio is not None:
             img = crop_to_ratio(img, crop_ratio)
         if size is not None: