{"cells":[{"metadata":{},"cell_type":"markdown","source":"# About this kernel\n\nImplement 4 point transformation.\n\n* Reference: [Finding the Transform matrix from 4 projected points (with Javascript)](https://math.stackexchange.com/a/339033/33138)\n* Both single image and batch images operations are implemented.\n* For random transformations, I set some limits on the 4 target points to avoid too distruptive transformations. You can change the values inside `random_4_points_2D()` or `random_4_points_2D_batch()` if you want.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import tensorflow as tf\nimport numpy as np\nfrom matplotlib import pyplot as plt\nimport re\nimport math\nimport datetime\nfrom kaggle_datasets import KaggleDatasets\nimport tensorflow.keras.backend as K\nAUTO = tf.data.experimental.AUTOTUNE","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"M = tf.convert_to_tensor([[1, 2, 3], [3, 4, 5], [4, 5, 6]], dtype=tf.float64)\nmask = tf.convert_to_tensor([[0, 0, 0], [0, 1, 1], [0, 1, 1]])\nindices = tf.convert_to_tensor(\n    [\n        [[1, 1], [1, 2]],\n        [[2, 1], [2, 2]]\n    ]\n)\ntf.gather_nd(M, indices)\nindices = tf.convert_to_tensor(\n    [\n        [\n            [[1, 1], [1, 2]],\n            [[2, 1], [2, 2]]\n        ],\n        [\n            [[1, 0], [1, 2]],\n            [[2, 0], [2, 2]]\n        ],\n        [\n            [[1, 0], [1, 1]],\n            [[2, 0], [2, 1]]\n        ],\n        [\n            [[0, 1], [0, 2]],\n            [[2, 1], [2, 2]]\n        ],\n        [\n            [[0, 0], [0, 2]],\n            [[2, 0], [2, 2]]\n        ],\n        [\n            [[0, 0], [0, 1]],\n            [[2, 0], [2, 1]]\n        ],\n        [\n            [[0, 1], [0, 2]],\n            [[1, 1], [1, 2]]\n        ],\n        [\n            [[0, 0], [0, 2]],\n            [[1, 0], [1, 2]]\n        ],\n        [\n            [[0, 0], [0, 1]],\n            [[1, 0], [1, 1]]\n        ]         \n    ]\n)\nindices = tf.reshape(indices, shape=(3, 3, 2, 2, 2))\nindices\n\ntf.gather_nd(M, indices)\n\ntf.linalg.det(tf.gather_nd(M, indices))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Common","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"image_size = 192\nIMAGE_SIZE = [image_size, image_size]\nBATCH_SIZE = 64\nAUG_BATCH = BATCH_SIZE\n\n# Data access\nGCS_DS_PATH = KaggleDatasets().get_gcs_path('tpu-getting-started')\n\n\nGCS_PATH_SELECT = { # available image sizes\n    192: GCS_DS_PATH + '/tfrecords-jpeg-192x192'\n\n}\nGCS_PATH = GCS_PATH_SELECT[IMAGE_SIZE[0]]\n\nTRAINING_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/train/*.tfrec')\n\ndef decode_image(image_data):\n    image = tf.image.decode_jpeg(image_data, channels=3)\n    image = tf.cast(image, tf.float32) / 255.0  # convert image to floats in [0, 1] range\n    image = tf.reshape(image, [*IMAGE_SIZE, 3]) # explicit size needed for TPU\n    return image\n\ndef read_labeled_tfrecord(example):\n    LABELED_TFREC_FORMAT = {\n        \"image\": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring\n        \"class\": tf.io.FixedLenFeature([], tf.int64),  # shape [] means single element\n    }\n    example = tf.io.parse_single_example(example, LABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    label = tf.cast(example['class'], tf.int32)\n    return image, label # returns a dataset of (image, label) pairs\n\ndef load_dataset(filenames, labeled = True, ordered = False):\n    # Read from TFRecords. For optimal performance, reading from multiple files at once and\n    # Diregarding data order. Order does not matter since we will be shuffling the data anyway\n    \n    ignore_order = tf.data.Options()\n    if not ordered:\n        ignore_order.experimental_deterministic = False # disable order, increase speed\n        \n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads = AUTO) # automatically interleaves reads from multiple files\n    dataset = dataset.with_options(ignore_order) # use data as soon as it streams in, rather than in its original order\n    dataset = dataset.map(read_labeled_tfrecord if labeled else read_unlabeled_tfrecord, num_parallel_calls = AUTO) # returns a dataset of (image, label) pairs if labeled = True or (image, id) pair if labeld = False\n    return dataset\n\ndef count_data_items(filenames):\n    # the number of data items is written in the name of the .tfrec files, i.e. flowers00-230.tfrec = 230 data items\n    n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) for filename in filenames]\n    return np.sum(n)\n\nNUM_TRAINING_IMAGES = int(count_data_items(TRAINING_FILENAMES))\n\nprint('Dataset: {} training images'.format(NUM_TRAINING_IMAGES))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def batch_to_numpy_images_and_labels(data):\n    \n    images = data\n    numpy_images = images.numpy()\n\n    # If no labels, only image IDs, return None for labels (this is the case for test data)\n    return numpy_images\n\n\ndef display_one_flower(image, title, subplot, red=False, titlesize=16):\n    plt.subplot(*subplot)\n    plt.axis('off')\n    plt.imshow(image)\n    if len(title) > 0:\n        plt.title(title, fontsize=int(titlesize) if not red else int(titlesize/1.2), color='red' if red else 'black', fontdict={'verticalalignment':'center'}, pad=int(titlesize/1.5))\n    return (subplot[0], subplot[1], subplot[2]+1)\n\n\ndef display_batch_of_images(databatch, predictions=None):\n    \n    \"\"\"This will work with:\n    display_batch_of_images(images)\n    \"\"\"\n    \n    # data\n    images = batch_to_numpy_images_and_labels(databatch)\n    labels = None\n    \n    if labels is None:\n        labels = [None for _ in enumerate(images)]\n        \n    # auto-squaring: this will drop data that does not fit into square or square-ish rectangle\n    rows = int(math.sqrt(len(images)))\n    cols = len(images)//rows\n        \n    # size and spacing\n    FIGSIZE = 13.0\n    SPACING = 0.1\n    subplot=(rows,cols,1)\n    if rows < cols:\n        plt.figure(figsize=(FIGSIZE,FIGSIZE/cols*rows))\n    else:\n        plt.figure(figsize=(FIGSIZE/rows*cols,FIGSIZE))\n    \n    # display\n    for i, (image, label) in enumerate(zip(images[:rows*cols], labels[:rows*cols])):\n        title = '' if label is None else CLASSES[label]\n        correct = True\n        if predictions is not None:\n            title, correct = title_from_label_and_target(predictions[i], label)\n        dynamic_titlesize = FIGSIZE*SPACING/max(rows,cols)*40+3 # magic formula tested to work from 1x1 to 10x10 images\n        subplot = display_one_flower(image, title, subplot, not correct, titlesize=dynamic_titlesize)\n    \n    #layout\n    plt.tight_layout()\n    if label is None and predictions is None:\n        plt.subplots_adjust(wspace=0, hspace=0)\n    else:\n        plt.subplots_adjust(wspace=SPACING, hspace=SPACING)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_training_dataset(dataset, batch_size=None, advanced_aug=True, repeat=True, with_labels=True, drop_remainder=False):\n    \n    if not with_labels:\n        dataset = dataset.map(lambda image, label: image, num_parallel_calls=AUTO)\n    \n    if advanced_aug:\n        dataset = dataset.map(transform, num_parallel_calls=AUTO)\n    \n    if type(repeat) == bool and repeat:\n        dataset = dataset.repeat() # the training dataset must repeat for several epochs\n    elif type(repeat) == int and repeat > 0:\n        dataset = dataset.repeat(repeat)\n    \n    dataset = dataset.shuffle(2048)\n    \n    if batch_size is None:\n        batch_size = BATCH_SIZE\n    \n    dataset = dataset.batch(batch_size, drop_remainder=drop_remainder)\n    dataset = dataset.prefetch(AUTO) # prefetch next batch while training (autotune prefetch buffer size)\n    \n    return dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = load_dataset(TRAINING_FILENAMES, labeled=True)\ntraining_dataset = get_training_dataset(dataset, advanced_aug=False, repeat=1, with_labels=True)\n\nimages, labels = next(iter(training_dataset.take(1)))\nprint(images.shape)\n\ndisplay_batch_of_images(images)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 4 point transformation","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def random_4_points_2D(height, width):\n    \"\"\"Generate 4 random 2-D points.\n    \n    The 4 points are inside a rectangle with the same center as the above rectangle but with side length being approximately 1.5 times.\n    This choice is to avoid the image being transformed too disruptively.\n    \n    Each point is created first by making it close to the corresponding corner points determined by the rectangle, i.e\n    [0, 0], [0, width], [height, width] and [height, 0] respectively. Then the 4 points are randomly shifted module 4.\n    \n    Args:\n        height: 0-D tensor, height of a reference rectangle.\n        width: 0-D tensor, width of a reference rectangle.\n        \n    Returns:\n        points: 2-D tensor of shape [4, 2]\n    \"\"\"\n\n    sy = height // 4\n    sx = width // 4\n        \n    h, w = height, width\n    \n    y1 = tf.random.uniform(minval = -sy, maxval = sy, shape=[], dtype=tf.int64)\n    x1 = tf.random.uniform(minval = -sx, maxval = sx, shape=[], dtype=tf.int64)\n    \n    y2 = tf.random.uniform(minval = -sy, maxval = sy, shape=[], dtype=tf.int64)\n    x2 = tf.random.uniform(minval = 3 * sx, maxval = 5 * sx, shape=[], dtype=tf.int64)\n\n    y3 = tf.random.uniform(minval = 3 * sy, maxval = 5 * sy, shape=[], dtype=tf.int64)\n    x3 = tf.random.uniform(minval = 3 * sx, maxval = 5 * sx, shape=[], dtype=tf.int64)    \n    \n    y4 = tf.random.uniform(minval = 3 * sy, maxval = 5 * sy, shape=[], dtype=tf.int64)\n    x4 = tf.random.uniform(minval = -sx, maxval = sx, shape=[], dtype=tf.int64)\n                \n    points = tf.convert_to_tensor([[y1, x1], [y2, x2], [y3, x3], [y4, x4], [y1, x1], [y2, x2], [y3, x3], [y4, x4]])\n    start_index = tf.random.uniform(minval=0, maxval=4, shape=[], dtype=tf.int64)\n    points = points[start_index: start_index + 4]\n        \n    return points\n\n\ndef random_4_point_transform_2D(image):\n    \"\"\"Apply 4 point transformation on 2-D image `image` with randomly generated 4 points on target spaces.\n    \n    On source space, the 4 points are the corner points, i.e [0, 0], [0, width], [height, width] and [height, 0].\n    On target space, the 4 points are randomly generated by `random_4_points_2D()`.\n    \"\"\"\n    \n    height, width = image.shape[:2]\n\n    # 4 corner points in source image\n    # shape = [4, 2]\n    src_pts = tf.convert_to_tensor([[0, 0], [0, width], [height, width], [height, 0]])\n\n    # 4 points in target image\n    # shape = [4, 2]\n    tgt_pts = random_4_points_2D(height, width)\n    \n    tgt_image = four_point_transform_2D(image, src_pts, tgt_pts)\n\n    return tgt_image\n\n\ndef four_point_transform_2D(image, src_pts, tgt_pts):\n    \"\"\"Apply 4 point transformation determined by `src_pts` and `tgt_pts` on 2-D image `image`.\n    \n    Args:\n        image: 2-D tensor of shape [height, width], or 3-D tensor of shape [height, width, channels]\n        src_pts: 2-D tensor of shape [4, 2]\n        tgt_pts: 2-D tensor of shape [4, 2]\n        \n    Returns:\n        A tensor with the same shape as `image`.\n    \"\"\"\n    \n    src_to_tgt_mat = get_src_to_tgt_mat_2D(src_pts, tgt_pts)\n    \n    tgt_image = transform_by_perspective_matrix_2D(image, src_to_tgt_mat)\n    \n    return tgt_image\n\n\ndef transform_by_perspective_matrix_2D(image, src_to_tgt_mat):\n    \"\"\"Transform a 2-D image by a prespective transformation matrix\n    \n    Args:\n        image: 2-D tensor of shape [height, width], or 3-D tensor of shape [height, width, channels]\n        src_to_tgt_mat: 2-D tensor of shape [3, 3]. This is the transformation matrix mapping the source space to the target space.\n        \n    Returns:\n        A tensor with the same shape as `image`.        \n    \"\"\"\n\n    height, width = image.shape[:2]\n\n    # shape = (3, 3)\n    tgt_to_src_mat = tf.linalg.inv(src_to_tgt_mat)\n        \n    # prepare y coordinates\n    # shape = [height * width]\n    ys = tf.repeat(tf.range(height), width)  \n\n    # prepare x coordinates\n    # shape = [height * width]\n    xs = tf.tile(tf.range(width), [height])\n\n    # prepare indices in target space\n    # shape = [2, height * width]\n    tgt_indices = tf.stack([ys, xs], axis=0)\n    \n    # Change to projective coordinates in the target space by adding ones\n    # shape = [3, height * width]\n    tgt_indices_homo = tf.concat([tgt_indices, tf.ones(shape=[1, height * width], dtype=tf.int32)], axis=0)\n    \n    # Get the corresponding projective coordinate in the source space\n    # shape = [3, height * width]\n    src_indices_homo = tf.linalg.matmul(tgt_to_src_mat, tf.cast(tgt_indices_homo, dtype=tf.float64))\n    \n    # normalize the projective coordinates\n    # shape = [3, height * width]\n    src_indices_normalized = src_indices_homo[:3, :] / src_indices_homo[2:, :]\n    \n    # Get the affine coordinate by removing ones\n    # shape = [2, height * width]\n    src_indices_affine = tf.cast(src_indices_normalized, dtype=tf.int64)[:2, :]\n    \n    # Mask the points outside the range\n    # shape = [height * width]\n    y_mask = tf.logical_and(src_indices_affine[0] >= 0, src_indices_affine[0] <= height - 1)\n    x_mask = tf.logical_and(src_indices_affine[1] >= 0, src_indices_affine[1] <= width - 1)\n    mask = tf.logical_and(y_mask, x_mask)\n    \n    # clip the coordinates\n    # shape = [2, height * width]\n    src_indices = tf.clip_by_value(src_indices_affine, clip_value_min=0, clip_value_max=[[height - 1], [width - 1]])\n    \n    # Get a collection of (y_coord, x_coord)\n    # shape = [height * width, 2]\n    src_indices = tf.transpose(src_indices)\n    \n    # shape = [height * width, channels]\n    tgt_image = tf.gather_nd(image, src_indices)\n    \n    # Set pixel to 0 by using the mask\n    tgt_image = tgt_image * tf.cast(mask[:, tf.newaxis], tf.float32)\n    \n    # reshape to [height, width, channels]\n    tgt_image = tf.reshape(tgt_image, image.shape)\n\n    return tgt_image\n\n\ndef get_src_to_tgt_mat_2D(src_pts, tgt_pts):\n    \"\"\"Get the perspective transformation matrix from the source space to the target space, which maps the 4 source points to the 4 target points.\n    \n    Args:\n        src_pts: 2-D tensor of shape [4, 2]\n        tgt_pts: 2-D tensor of shape [4, 2]\n        \n    Returns:\n        2-D tensor of shape [3, 3]\n    \"\"\"\n    \n    src_pts = tf.cast(src_pts, tf.int64)\n    tgt_pts = tf.cast(tgt_pts, tf.int64)\n    \n    # The perspective transformation matrix mapping basis vectors and (1, 1, 1) to `src_pts`\n    # shape = [3, 3]\n    src_mat = get_transformation_mat_2D(src_pts)\n    \n    # The perspective transformation matrix mapping basis vectors and (1, 1, 1) to `tgt_pts`\n    # shape = [3, 3]\n    tgt_mat = get_transformation_mat_2D(tgt_pts)\n    \n    # The perspective transformation matrix mapping `src_pts` to `tgt_pts`\n    # shape = [3, 3]\n    src_to_tgt_mat = tf.linalg.matmul(tgt_mat, tf.linalg.inv(src_mat))\n    \n    return src_to_tgt_mat\n  \n    \ndef get_transformation_mat_2D(four_pts):\n    \"\"\"Get the perspective transformation matrix from a space to another space, which maps the basis vectors and (1, 1, 1) to the 4 points defined by `four_pts`.\n    \n    Args:\n        four_pts: 2-D tensor of shape [4, 2]\n        \n    Returns:\n        2-D tensor of shape [3, 3]        \n    \"\"\"\n    \n    # Change to projective coordinates by adding ones\n    # shape = [3, 4]\n    #pts_homo = tf.concat([tf.transpose(four_pts), tf.ones(shape=[1, 4], dtype=tf.int64)], axis=0)\n    pts_homo = tf.transpose(tf.concat([four_pts, tf.ones(shape=[4, 1], dtype=tf.int64)], axis=-1))\n    \n    pts_homo = tf.cast(pts_homo, tf.float64)\n    \n    # Find `scalars` such that: src_pts_homo[:, 3:] * scalars == src_pts_homo[:, 3:]\n    # shape = [3, 3]\n    inv_mat = tf.linalg.inv(pts_homo[:, :3])\n    # shape = [3, 1]\n    scalars = tf.linalg.matmul(inv_mat, pts_homo[:, 3:])\n    \n    # Get the matrix transforming unit vectors to the 4 source points\n    # shape = [3, 3]\n    mat = tf.transpose(tf.transpose(pts_homo[:, :3]) * scalars)\n    \n    return mat","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"new_image = random_4_point_transform_2D(images[0])\ndisplay_batch_of_images(tf.convert_to_tensor([images[0], new_image]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"new_image = random_4_point_transform_2D(images[1])\ndisplay_batch_of_images(tf.convert_to_tensor([images[1], new_image]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"new_image = random_4_point_transform_2D(images[2])\ndisplay_batch_of_images(tf.convert_to_tensor([images[2], new_image]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"new_image = random_4_point_transform_2D(images[3])\ndisplay_batch_of_images(tf.convert_to_tensor([images[3], new_image]))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Batch implementation","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def random_4_points_2D_batch(height, width, batch_size):\n    \"\"\"Generate `batch_size * 4` random 2-D points.\n    \n    Each 4 points are inside a rectangle with the same center as the above rectangle but with side length being approximately 1.5 times.\n    This choice is to avoid the image being transformed too disruptively.\n\n    Each point is created first by making it close to the corresponding corner points determined by the rectangle, i.e\n    [0, 0], [0, width], [height, width] and [height, 0] respectively. Then the 4 points are randomly shifted module 4.\n    \n    Args:\n        height: 0-D tensor, height of a reference rectangle.\n        width: 0-D tensor, width of a reference rectangle.\n        batch_size: 0-D tensor, the number of 4 points to be generated\n        \n    Returns:\n        points: 3-D tensor of shape [batch_size, 4, 2]\n    \"\"\"\n\n    sy = height // 4\n    sx = width // 4\n        \n    h, w = height, width\n    \n    y1 = tf.random.uniform(minval = -sy, maxval = sy, shape=[batch_size], dtype=tf.int64)\n    x1 = tf.random.uniform(minval = -sx, maxval = sx, shape=[batch_size], dtype=tf.int64)\n    \n    y2 = tf.random.uniform(minval = -sy, maxval = sy, shape=[batch_size], dtype=tf.int64)\n    x2 = tf.random.uniform(minval = 3 * sx, maxval = 5 * sx, shape=[batch_size], dtype=tf.int64)\n\n    y3 = tf.random.uniform(minval = 3 * sy, maxval = 5 * sy, shape=[batch_size], dtype=tf.int64)\n    x3 = tf.random.uniform(minval = 3 * sx, maxval = 5 * sx, shape=[batch_size], dtype=tf.int64)    \n    \n    y4 = tf.random.uniform(minval = 3 * sy, maxval = 5 * sy, shape=[batch_size], dtype=tf.int64)\n    x4 = tf.random.uniform(minval = -sx, maxval = sx, shape=[batch_size], dtype=tf.int64)\n            \n    # shape = [4, 2, batch_size]\n    points = tf.convert_to_tensor([[y1, x1], [y2, x2], [y3, x3], [y4, x4], [y1, x1], [y2, x2], [y3, x3], [y4, x4]])\n    \n    # shape = [batch_size, 4, 2]\n    points = tf.transpose(points, perm=[2, 0, 1])\n    \n    # Trick to get random rotation\n    # shape = [batch_size, 8, 2]\n    points = tf.tile(points, multiples=[1, 2, 1])    \n    # shape = [batch_size]\n    start_indices = tf.random.uniform(minval=0, maxval=4, shape=[batch_size], dtype=tf.int64)\n    # shape = [batch_size, 4]\n    indices = start_indices[:, tf.newaxis] + tf.range(4, dtype=tf.int64)[tf.newaxis, :]\n    # shape = [batch_size, 4, 2]\n    indices = tf.stack([tf.broadcast_to(tf.range(batch_size, dtype=tf.int64)[:, tf.newaxis], shape=[batch_size, 4]), indices], axis=2)    \n    \n    # shape = [batch_size, 4, 2]\n    points = tf.gather_nd(points, indices)\n        \n    return points\n\n\ndef random_4_point_transform_2D_batch(images):\n    \"\"\"Apply 4 point transformation on 2-D images `images` with randomly generated 4 points on target spaces.\n    \n    On source space, the 4 points are the corner points, i.e [0, 0], [0, width], [height, width] and [height, 0].\n    On target space, the 4 points are randomly generated by `random_4_points_2D_batch()`.\n    \"\"\"\n\n    batch_size, height, width = images.shape[:3]\n\n    # 4 corner points in source image\n    # shape = [batch_size, 4, 2]\n    src_pts = tf.convert_to_tensor([[0, 0], [0, width], [height, width], [height, 0]])\n    src_pts = tf.broadcast_to(src_pts, shape=[batch_size, 4, 2])\n\n    # 4 points in target image\n    # shape = [batch_size, 4, 2]\n    tgt_pts = random_4_points_2D_batch(height, width, batch_size)\n    \n    tgt_images = four_point_transform_2D_batch(images, src_pts, tgt_pts)\n\n    return tgt_images\n\n\ndef four_point_transform_2D_batch(images, src_pts, tgt_pts):\n    \"\"\"Apply 4 point transformation determined by `src_pts` and `tgt_pts` on 2-D images `images`.\n    \n    Args:\n        images: 3-D tensor of shape [batch_size, height, width], or 4-D tensor of shape [batch_size, height, width, channels]\n        src_pts: 3-D tensor of shape [batch_size, 4, 2]\n        tgt_pts: 3-D tensor of shape [batch_size, 4, 2]\n        \n    Returns:\n        A tensor with the same shape as `images`.\n    \"\"\"\n    \n    src_to_tgt_mat = get_src_to_tgt_mat_2D_batch(src_pts, tgt_pts)\n    \n    tgt_images = transform_by_perspective_matrix_2D_batch(images, src_to_tgt_mat)\n    \n    return tgt_images\n\n\ndef transform_by_perspective_matrix_2D_batch(images, src_to_tgt_mat):\n    \"\"\"Transform 2-D images by prespective transformation matrices\n    \n    Args:\n        images: 3-D tensor of shape [batch_size, height, width], or 4-D tensor of shape [batch_size, height, width, channels]\n        src_to_tgt_mat: 3-D tensor of shape [batch_size, 3, 3]. This is the transformation matrix mapping the source space to the target space.\n        \n    Returns:\n        A tensor with the same shape as `image`.        \n    \"\"\"\n\n    batch_size, height, width = images.shape[:3]\n\n    # shape = (3, 3)\n    tgt_to_src_mat = tf.linalg.inv(src_to_tgt_mat)\n        \n    # prepare y coordinates\n    # shape = [height * width]\n    ys = tf.repeat(tf.range(height), width) \n    \n    # prepare x coordinates\n    # shape = [height * width]\n    xs = tf.tile(tf.range(width), [height])\n\n    # prepare indices in target space\n    # shape = [2, height * width]\n    tgt_indices = tf.stack([ys, xs], axis=0)\n    \n    # Change to projective coordinates in the target space by adding ones\n    # shape = [3, height * width]\n    tgt_indices_homo = tf.concat([tgt_indices, tf.ones(shape=[1, height * width], dtype=tf.int32)], axis=0)\n    \n    # Get the corresponding projective coordinate in the source space\n    # shape = [batch_size, 3, height * width]\n    src_indices_homo = tf.linalg.matmul(tgt_to_src_mat, tf.cast(tgt_indices_homo, dtype=tf.float64))\n    \n    # normalize the projective coordinates\n    # shape = [batch_size, 3, height * width]\n    src_indices_normalized = src_indices_homo[:, :3, :] / src_indices_homo[:, 2:, :]\n    \n    # Get the affine coordinate by removing ones\n    # shape = [batch_size, 2, height * width]\n    src_indices_affine = tf.cast(src_indices_normalized, dtype=tf.int64)[:, :2, :]\n    \n    # Mask the points outside the range\n    # shape = [batch_size, height * width]\n    y_mask = tf.logical_and(src_indices_affine[:, 0] >= 0, src_indices_affine[:, 0] <= height - 1)\n    x_mask = tf.logical_and(src_indices_affine[:, 1] >= 0, src_indices_affine[:, 1] <= width - 1)\n    mask = tf.logical_and(y_mask, x_mask)\n    \n    # clip the coordinates\n    # shape = [batch_size, 2, height * width]\n    src_indices = tf.clip_by_value(src_indices_affine, clip_value_min=0, clip_value_max=[[height - 1], [width - 1]])\n    \n    # Get a collection of (y_coord, x_coord)\n    # shape = [batch_size, height * width, 2]\n    src_indices = tf.transpose(src_indices, perm=[0, 2, 1])\n    \n    # shape = [batch_size, height * width, channels]\n    tgt_images = tf.gather_nd(images, src_indices, batch_dims=1)\n    \n    # Set pixel to 0 by using the mask\n    tgt_images = tgt_images * tf.cast(mask[:, :, tf.newaxis], tf.float32)\n    \n    # reshape to [height, width, channels]\n    tgt_images = tf.reshape(tgt_images, images.shape)\n\n    return tgt_images\n\n\ndef get_src_to_tgt_mat_2D_batch(src_pts, tgt_pts):\n    \"\"\"Get the perspective transformation matrix from the source space to the target space, which maps the 4 source points to the 4 target points.\n    \n    Args:\n        src_pts: 3-D tensor of shape [batch_size, 4, 2]\n        tgt_pts: 3-D tensor of shape [batch_size, 4, 2]\n        \n    Returns:\n        2-D tensor of shape [batch_size, 3, 3]\n    \"\"\"\n    \n    src_pts = tf.cast(src_pts, tf.int64)\n    tgt_pts = tf.cast(tgt_pts, tf.int64)\n    \n    # The perspective transformation matrix mapping basis vectors and (1, 1, 1) to `src_pts`\n    # shape = [batch_size, 3, 3]\n    src_mat = get_transformation_mat_2D_batch(src_pts)\n    \n    # The perspective transformation matrix mapping basis vectors and (1, 1, 1) to `tgt_pts`\n    # shape = [3, 3]\n    tgt_mat = get_transformation_mat_2D_batch(tgt_pts)\n    \n    # The perspective transformation matrix mapping `src_pts` to `tgt_pts`\n    # shape = [3, 3]\n    src_to_tgt_mat = tf.linalg.matmul(tgt_mat, tf.linalg.inv(src_mat))\n    \n    return src_to_tgt_mat\n  \n    \ndef get_transformation_mat_2D_batch(four_pts):\n    \"\"\"Get the perspective transformation matrix from a space to another space, which maps the basis vectors and (1, 1, 1) to the 4 points defined by `four_pts`.\n    \n    Args:\n        four_pts: 3-D tensor of shape [batch_size, 4, 2]\n        \n    Returns:\n        3-D tensor of shape [batch_size, 3, 3]        \n    \"\"\"\n    \n    batch_size = four_pts.shape[0]\n    \n    # Change to projective coordinates by adding ones\n    # shape = [batch_size, 3, 4]\n    pts_homo = tf.transpose(tf.concat([four_pts, tf.ones(shape=[batch_size, 4, 1], dtype=tf.int64)], axis=-1), perm=[0, 2, 1])\n    \n    pts_homo = tf.cast(pts_homo, tf.float64)\n    \n    # Find `scalars` such that: src_pts_homo[:, 3:] * scalars == src_pts_homo[:, 3:]\n    # shape = [batch_size 3, 3]\n    inv_mat = tf.linalg.inv(pts_homo[:, :, :3])\n    # shape = [batch_size, 3, 1]\n    scalars = tf.linalg.matmul(inv_mat, pts_homo[:, :, 3:])\n    \n    # Get the matrix transforming unit vectors to the 4 source points\n    # shape = [batch_size, 3, 3]    \n    mat = tf.transpose(tf.transpose(pts_homo[:, :, :3], perm=[0, 2, 1]) * scalars, perm=[0, 2, 1])\n    \n    return mat","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"new_images = random_4_point_transform_2D_batch(tf.repeat(images[:16], axis=0, repeats=4))\ndisplay_batch_of_images(new_images)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"new_images = random_4_point_transform_2D_batch(tf.repeat(images[16:32], axis=0, repeats=4))\ndisplay_batch_of_images(new_images)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"new_images = random_4_point_transform_2D_batch(images)\ndisplay_batch_of_images(new_images)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"n_iter = 1000\n\nstart = datetime.datetime.now()\nfor i in range(n_iter):\n    random_4_point_transform_2D_batch(images) \nend = datetime.datetime.now()\ntiming = (end - start).total_seconds() / n_iter\nprint(f\"batch_4_pt_transformation: {timing}\")","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}