SparseImageWarp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TensorFlow's `sparse_image_warp` in Pytorch\n",
    "\n",
    "This notebook is an implementation of TensorFlow's [sparse_image_warp](https://www.tensorflow.org/versions/r1.9/api_docs/python/tf/contrib/image/sparse_image_warp?hl=en).\n",
    "\n",
    "We call `sparse_image_warp` in the accompanying SpecAugment notebook for use in the Time Warp data augmentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # NOTICE: YOU MUST FIRST CALL `%store spectro` in the other notebook to use the spectro value here\n",
    "# %store -r spectro"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Export\n",
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our function expects inputs of the following shape:\n",
    "\n",
    "```py\n",
    "# source_control_point_locations: `[batch, num_control_points, 2]` float\n",
    "#   `Tensor`\n",
    "# dest_control_point_locations: `[batch, num_control_points, 2]` float\n",
    "#   `Tensor`\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Export\n",
    "def sparse_image_warp(img_tensor,\n",
    "                      source_control_point_locations,\n",
    "                      dest_control_point_locations,\n",
    "                      interpolation_order=2,\n",
    "                      regularization_weight=0.0,\n",
    "                      num_boundaries_points=0):\n",
    "    device = img_tensor.device\n",
    "    control_point_flows = (dest_control_point_locations - source_control_point_locations)   \n",
    "    \n",
    "#     clamp_boundaries = num_boundary_points > 0\n",
    "#     boundary_points_per_edge = num_boundary_points - 1\n",
    "    batch_size, image_height, image_width = img_tensor.shape\n",
    "    flattened_grid_locations = get_flat_grid_locations(image_height, image_width, device)\n",
    "\n",
    "    # IGNORED FOR OUR BASIC VERSION...\n",
    "#     flattened_grid_locations = constant_op.constant(\n",
    "#         _expand_to_minibatch(flattened_grid_locations, batch_size), image.dtype)\n",
    "\n",
    "#     if clamp_boundaries:\n",
    "#       (dest_control_point_locations,\n",
    "#        control_point_flows) = _add_zero_flow_controls_at_boundary(\n",
    "#            dest_control_point_locations, control_point_flows, image_height,\n",
    "#            image_width, boundary_points_per_edge)\n",
    "\n",
    "    flattened_flows = interpolate_spline(\n",
    "        dest_control_point_locations,\n",
    "        control_point_flows,\n",
    "        flattened_grid_locations,\n",
    "        interpolation_order,\n",
    "        regularization_weight)\n",
    "\n",
    "    dense_flows = create_dense_flows(flattened_flows, batch_size, image_height, image_width)\n",
    "\n",
    "    warped_image = dense_image_warp(img_tensor, dense_flows)\n",
    "\n",
    "    return warped_image, dense_flows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Export\n",
    "def get_grid_locations(image_height, image_width, device):\n",
    "    y_range = torch.linspace(0, image_height - 1, image_height, device=device)\n",
    "    x_range = torch.linspace(0, image_width - 1, image_width, device=device)\n",
    "    y_grid, x_grid = torch.meshgrid(y_range, x_range)\n",
    "    return torch.stack((y_grid, x_grid), -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Export\n",
    "def flatten_grid_locations(grid_locations, image_height, image_width):\n",
    "    return torch.reshape(grid_locations, [image_height * image_width, 2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Export\n",
    "def get_flat_grid_locations(image_height, image_width, device):\n",
    "    y_range = torch.linspace(0, image_height - 1, image_height, device=device)\n",
    "    x_range = torch.linspace(0, image_width - 1, image_width, device=device)\n",
    "    y_grid, x_grid = torch.meshgrid(y_range, x_range)\n",
    "    return torch.stack((y_grid, x_grid), -1).reshape([image_height * image_width, 2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Export\n",
    "def create_dense_flows(flattened_flows, batch_size, image_height, image_width):\n",
    "    # possibly .view\n",
    "    return torch.reshape(flattened_flows, [batch_size, image_height, image_width, 2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interpolate Spline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Export\n",
    "def interpolate_spline(train_points, train_values, query_points, order, regularization_weight=0.0,):\n",
    "    # First, fit the spline to the observed data.\n",
    "    w, v = solve_interpolation(train_points, train_values, order, regularization_weight)\n",
    "    # Then, evaluate the spline at the query locations.\n",
    "    query_values = apply_interpolation(query_points, train_points, w, v, order)\n",
    "\n",
    "    return query_values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Export\n",
    "def solve_interpolation(train_points, train_values, order, regularization_weight, eps=1e-7):\n",
    "    device = train_points.device\n",
    "    b, n, d = train_points.shape\n",
    "    k = train_values.shape[-1]\n",
    "\n",
    "    # First, rename variables so that the notation (c, f, w, v, A, B, etc.)\n",
    "    # follows https://en.wikipedia.org/wiki/Polyharmonic_spline.\n",
    "    # To account for python style guidelines we use\n",
    "    # matrix_a for A and matrix_b for B.\n",
    "    \n",
    "    c = train_points\n",
    "    f = train_values.float()\n",
    "    \n",
    "    matrix_a = phi(cross_squared_distance_matrix(c,c), order).unsqueeze(0)  # [b, n, n]\n",
    "#     if regularization_weight > 0:\n",
    "#         batch_identity_matrix = array_ops.expand_dims(\n",
    "#           linalg_ops.eye(n, dtype=c.dtype), 0)\n",
    "#         matrix_a += regularization_weight * batch_identity_matrix\n",
    "\n",
    "    # Append ones to the feature values for the bias term in the linear model.\n",
    "    ones = torch.ones(n, dtype=train_points.dtype, device=device).view([-1, n, 1])\n",
    "    matrix_b = torch.cat((c, ones), 2).float()  # [b, n, d + 1]\n",
    "\n",
    "    # [b, n + d + 1, n]\n",
    "    left_block = torch.cat((matrix_a, torch.transpose(matrix_b, 2, 1)), 1)\n",
    "\n",
    "    num_b_cols = matrix_b.shape[2]  # d + 1\n",
    "\n",
    "    # In Tensorflow, zeros are used here. Pytorch solve fails with zeros for some reason we don't understand.\n",
    "    # So instead we use very tiny randn values (variance of one, zero mean) on one side of our multiplication.\n",
    "    lhs_zeros = torch.randn((b, num_b_cols, num_b_cols), device=device) *eps\n",
    "    right_block = torch.cat((matrix_b, lhs_zeros),\n",
    "                                   1)  # [b, n + d + 1, d + 1]\n",
    "    lhs = torch.cat((left_block, right_block),\n",
    "                           2)  # [b, n + d + 1, n + d + 1]\n",
    "\n",
    "    rhs_zeros = torch.zeros((b, d + 1, k), dtype=train_points.dtype, device=device).float()\n",
    "    rhs = torch.cat((f, rhs_zeros), 1)  # [b, n + d + 1, k]\n",
    "\n",
    "    # Then, solve the linear system and unpack the results.\n",
    "    X, LU = torch.solve(rhs, lhs)\n",
    "    w = X[:, :n, :]\n",
    "    v = X[:, n:, :]\n",
    "    return w, v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Export\n",
    "def cross_squared_distance_matrix(x, y):\n",
    "    \"\"\"Pairwise squared distance between two (batch) matrices' rows (2nd dim).\n",
    "        Computes the pairwise distances between rows of x and rows of y\n",
    "        Args:\n",
    "        x: [batch_size, n, d] float `Tensor`\n",
    "        y: [batch_size, m, d] float `Tensor`\n",
    "        Returns:\n",
    "        squared_dists: [batch_size, n, m] float `Tensor`, where\n",
    "        squared_dists[b,i,j] = ||x[b,i,:] - y[b,j,:]||^2\n",
    "    \"\"\"\n",
    "    x_norm_squared = torch.sum(torch.mul(x, x))\n",
    "    y_norm_squared = torch.sum(torch.mul(y, y))\n",
    "\n",
    "    x_y_transpose = torch.matmul(x.squeeze(0), y.squeeze(0).transpose(0,1))\n",
    "    \n",
    "    # squared_dists[b,i,j] = ||x_bi - y_bj||^2 = x_bi'x_bi- 2x_bi'x_bj + x_bj'x_bj\n",
    "    squared_dists = x_norm_squared - 2 * x_y_transpose + y_norm_squared\n",
    "\n",
    "    return squared_dists.float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Export\n",
    "def phi(r, order):\n",
    "    \"\"\"Coordinate-wise nonlinearity used to define the order of the interpolation.\n",
    "    See https://en.wikipedia.org/wiki/Polyharmonic_spline for the definition.\n",
    "    Args:\n",
    "    r: input op\n",
    "    order: interpolation order\n",
    "    Returns:\n",
    "    phi_k evaluated coordinate-wise on r, for k = r\n",
    "    \"\"\"\n",
    "    EPSILON=torch.tensor(1e-10, device=r.device)\n",
    "    # using EPSILON prevents log(0), sqrt0), etc.\n",
    "    # sqrt(0) is well-defined, but its gradient is not\n",
    "    if order == 1:\n",
    "        r = torch.max(r, EPSILON)\n",
    "        r = torch.sqrt(r)\n",
    "        return r\n",
    "    elif order == 2:\n",
    "        return 0.5 * r * torch.log(torch.max(r, EPSILON))\n",
    "    elif order == 4:\n",
    "        return 0.5 * torch.square(r) * torch.log(torch.max(r, EPSILON))\n",
    "    elif order % 2 == 0:\n",
    "        r = torch.max(r, EPSILON)\n",
    "        return 0.5 * torch.pow(r, 0.5 * order) * torch.log(r)\n",
    "    else:\n",
    "        r = torch.max(r, EPSILON)\n",
    "        return torch.pow(r, 0.5 * order)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Export\n",
    "def apply_interpolation(query_points, train_points, w, v, order):\n",
    "    \"\"\"Apply polyharmonic interpolation model to data.\n",
    "    Given coefficients w and v for the interpolation model, we evaluate\n",
    "    interpolated function values at query_points.\n",
    "    Args:\n",
    "    query_points: `[b, m, d]` x values to evaluate the interpolation at\n",
    "    train_points: `[b, n, d]` x values that act as the interpolation centers\n",
    "                    ( the c variables in the wikipedia article)\n",
    "    w: `[b, n, k]` weights on each interpolation center\n",
    "    v: `[b, d, k]` weights on each input dimension\n",
    "    order: order of the interpolation\n",
    "    Returns:\n",
    "    Polyharmonic interpolation evaluated at points defined in query_points.\n",
    "    \"\"\"\n",
    "    query_points = query_points.unsqueeze(0)\n",
    "    # First, compute the contribution from the rbf term.\n",
    "    pairwise_dists = cross_squared_distance_matrix(query_points.float(), train_points.float())\n",
    "    phi_pairwise_dists = phi(pairwise_dists, order)\n",
    "\n",
    "    rbf_term = torch.matmul(phi_pairwise_dists, w)\n",
    "\n",
    "    # Then, compute the contribution from the linear term.\n",
    "    # Pad query_points with ones, for the bias term in the linear model.\n",
    "    ones = torch.ones_like(query_points[..., :1])\n",
    "    query_points_pad = torch.cat((\n",
    "      query_points,\n",
    "      ones\n",
    "    ), 2).float()\n",
    "    linear_term = torch.matmul(query_points_pad, v)\n",
    "\n",
    "    return rbf_term + linear_term\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 4, 4]) torch.Size([1, 4, 2])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 4, 2]),\n",
       " tensor([[[ 6.0000,  0.0000,  1.0000,  0.0000],\n",
       "          [ 0.1667,  1.0000, -0.1667,  0.0000],\n",
       "          [ 0.1667,  1.0000,  6.0000,  1.0000],\n",
       "          [ 0.1667,  0.0000, -0.0278,  1.0278]]]))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_A = torch.tensor([[[1., 1., 6., 1.],\n",
    "         [1., 1., 0., 0.],\n",
    "         [6., 0., 1., 0.],\n",
    "         [1., 0., 0., 1.]]])\n",
    "test_B = torch.tensor([[[-1.,  3.],\n",
    "         [ 0.,  0.],\n",
    "         [ 0.,  0.],\n",
    "         [ 0.,  0.]]])\n",
    "print(test_A.shape, test_B.shape)\n",
    "X, LU = torch.solve(test_B, test_A)\n",
    "X.shape, LU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Export\n",
    "def dense_image_warp(image, flow):\n",
    "    \"\"\"Image warping using per-pixel flow vectors.\n",
    "    Apply a non-linear warp to the image, where the warp is specified by a dense\n",
    "    flow field of offset vectors that define the correspondences of pixel values\n",
    "    in the output image back to locations in the  source image. Specifically, the\n",
    "    pixel value at output[b, j, i, c] is\n",
    "    images[b, j - flow[b, j, i, 0], i - flow[b, j, i, 1], c].\n",
    "    The locations specified by this formula do not necessarily map to an int\n",
    "    index. Therefore, the pixel value is obtained by bilinear\n",
    "    interpolation of the 4 nearest pixels around\n",
    "    (b, j - flow[b, j, i, 0], i - flow[b, j, i, 1]). For locations outside\n",
    "    of the image, we use the nearest pixel values at the image boundary.\n",
    "    Args:\n",
    "    image: 4-D float `Tensor` with shape `[batch, height, width, channels]`.\n",
    "    flow: A 4-D float `Tensor` with shape `[batch, height, width, 2]`.\n",
    "    name: A name for the operation (optional).\n",
    "    Note that image and flow can be of type tf.half, tf.float32, or tf.float64,\n",
    "    and do not necessarily have to be the same type.\n",
    "    Returns:\n",
    "    A 4-D float `Tensor` with shape`[batch, height, width, channels]`\n",
    "    and same type as input image.\n",
    "    Raises:\n",
    "    ValueError: if height < 2 or width < 2 or the inputs have the wrong number\n",
    "    of dimensions.\n",
    "    \"\"\"\n",
    "    image = image.unsqueeze(3) # add a single channel dimension to image tensor\n",
    "    batch_size, height, width, channels = image.shape\n",
    "    device = image.device\n",
    "\n",
    "    # The flow is defined on the image grid. Turn the flow into a list of query\n",
    "    # points in the grid space.\n",
    "    grid_x, grid_y = torch.meshgrid(\n",
    "        torch.arange(width, device=device), torch.arange(height, device=device))\n",
    "    \n",
    "    stacked_grid = torch.stack((grid_y, grid_x), dim=2).float()\n",
    "    \n",
    "    batched_grid = stacked_grid.unsqueeze(-1).permute(3, 1, 0, 2)\n",
    "    \n",
    "    query_points_on_grid = batched_grid - flow\n",
    "    query_points_flattened = torch.reshape(query_points_on_grid,\n",
    "                                               [batch_size, height * width, 2])\n",
    "    # Compute values at the query points, then reshape the result back to the\n",
    "    # image grid.\n",
    "    interpolated = interpolate_bilinear(image, query_points_flattened)\n",
    "    interpolated = torch.reshape(interpolated,\n",
    "                                     [batch_size, height, width, channels])\n",
    "    return interpolated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Export\n",
    "def interpolate_bilinear(grid,\n",
    "                         query_points,\n",
    "                         name='interpolate_bilinear',\n",
    "                         indexing='ij'):\n",
    "    \"\"\"Similar to Matlab's interp2 function.\n",
    "    Finds values for query points on a grid using bilinear interpolation.\n",
    "    Args:\n",
    "    grid: a 4-D float `Tensor` of shape `[batch, height, width, channels]`.\n",
    "    query_points: a 3-D float `Tensor` of N points with shape `[batch, N, 2]`.\n",
    "    name: a name for the operation (optional).\n",
    "    indexing: whether the query points are specified as row and column (ij),\n",
    "      or Cartesian coordinates (xy).\n",
    "    Returns:\n",
    "    values: a 3-D `Tensor` with shape `[batch, N, channels]`\n",
    "    Raises:\n",
    "    ValueError: if the indexing mode is invalid, or if the shape of the inputs\n",
    "      invalid.\n",
    "    \"\"\"\n",
    "    if indexing != 'ij' and indexing != 'xy':\n",
    "        raise ValueError('Indexing mode must be \\'ij\\' or \\'xy\\'')\n",
    "\n",
    "\n",
    "    shape = grid.shape\n",
    "    if len(shape) != 4:\n",
    "      msg = 'Grid must be 4 dimensional. Received size: '\n",
    "      raise ValueError(msg + str(grid.shape))\n",
    "\n",
    "    batch_size, height, width, channels = grid.shape\n",
    "\n",
    "    shape = [batch_size, height, width, channels]\n",
    "    query_type = query_points.dtype\n",
    "    grid_type = grid.dtype\n",
    "    grid_device = grid.device\n",
    "\n",
    "    num_queries = query_points.shape[1]\n",
    "\n",
    "    alphas = []\n",
    "    floors = []\n",
    "    ceils = []\n",
    "    index_order = [0, 1] if indexing == 'ij' else [1, 0]\n",
    "    unstacked_query_points = query_points.unbind(2)\n",
    "\n",
    "    for dim in index_order:\n",
    "        queries = unstacked_query_points[dim]\n",
    "\n",
    "        size_in_indexing_dimension = shape[dim + 1]\n",
    "\n",
    "        # max_floor is size_in_indexing_dimension - 2 so that max_floor + 1\n",
    "        # is still a valid index into the grid.\n",
    "        max_floor = torch.tensor(size_in_indexing_dimension - 2, dtype=query_type, device=grid_device)\n",
    "        min_floor = torch.tensor(0.0, dtype=query_type, device=grid_device)\n",
    "        maxx = torch.max(min_floor, torch.floor(queries))\n",
    "        floor = torch.min(maxx, max_floor)\n",
    "        int_floor = floor.long()\n",
    "        floors.append(int_floor)\n",
    "        ceil = int_floor + 1\n",
    "        ceils.append(ceil)\n",
    "\n",
    "        # alpha has the same type as the grid, as we will directly use alpha\n",
    "        # when taking linear combinations of pixel values from the image.\n",
    "        \n",
    "        \n",
    "        alpha = (queries - floor).clone().detach().type(grid_type)\n",
    "        min_alpha = torch.tensor(0.0, dtype=grid_type, device=grid_device)\n",
    "        max_alpha = torch.tensor(1.0, dtype=grid_type, device=grid_device)\n",
    "        alpha = torch.min(torch.max(min_alpha, alpha), max_alpha)\n",
    "\n",
    "        # Expand alpha to [b, n, 1] so we can use broadcasting\n",
    "        # (since the alpha values don't depend on the channel).\n",
    "        alpha = torch.unsqueeze(alpha, 2)\n",
    "        alphas.append(alpha)\n",
    "\n",
    "    flattened_grid = torch.reshape(\n",
    "      grid, [batch_size * height * width, channels])\n",
    "    batch_offsets = torch.reshape(\n",
    "      torch.arange(batch_size, device=grid_device) * height * width, [batch_size, 1])\n",
    "\n",
    "    # This wraps array_ops.gather. We reshape the image data such that the\n",
    "    # batch, y, and x coordinates are pulled into the first dimension.\n",
    "    # Then we gather. Finally, we reshape the output back. It's possible this\n",
    "    # code would be made simpler by using array_ops.gather_nd.\n",
    "    def gather(y_coords, x_coords, name):\n",
    "        linear_coordinates = batch_offsets + y_coords * width + x_coords\n",
    "        gathered_values = torch.gather(flattened_grid.t(), 1, linear_coordinates)\n",
    "        return torch.reshape(gathered_values,\n",
    "                                 [batch_size, num_queries, channels])\n",
    "\n",
    "    # grab the pixel values in the 4 corners around each query point\n",
    "    top_left = gather(floors[0], floors[1], 'top_left')\n",
    "    top_right = gather(floors[0], ceils[1], 'top_right')\n",
    "    bottom_left = gather(ceils[0], floors[1], 'bottom_left')\n",
    "    bottom_right = gather(ceils[0], ceils[1], 'bottom_right')\n",
    "\n",
    "    interp_top = alphas[1] * (top_right - top_left) + top_left\n",
    "    interp_bottom = alphas[1] * (bottom_right - bottom_left) + bottom_left\n",
    "    interp = alphas[0] * (interp_bottom - interp_top) + interp_top\n",
    "\n",
    "    return interp\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}