Torch equivalent of skimage.transform.PiecewiseAffineTransform()

username: RuoyuGuo
cooked: <p>Hi, all. I am trying to reproduce a random warping algorithm in <a href="https://openaccess.thecvf.com/content/CVPR2022/supplemental/Zeng_SketchEdit_Mask-Free_Local_CVPR_2022_supplemental.pdf" rel="noopener nofollow ugc">SketchEdit: Mask-Free Local Image Manipulation with Partial Sketches supplemental</a>. I also would like to apply it as a pre-processing method to train my network.<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/7/b/7b9eafbb101646af73ead66cb901e65e32fdf5f7.png" alt="image" data-base62-sha1="hDAIcONmhcylZQJ6kxKDrJ56xYH" width="410" height="204"></p>
<p><strong>Original Algorithm</strong>: As shown in Figure 1.  They randomly sample some control points (blue points in Figure 1 (b))  and construct a triangular mesh using the Delaunay triangulation method.  To warp the<br>
modification area, we randomly move the control points by a small amount (∆x, ∆y) and apply the affine transformation on each triangle associated the new location.</p>
<p><strong>My solution with skimage</strong>:</p>
<pre><code class="lang-auto">img = READ_YOUR_IMAGE
image = resize(img, (256,256))

SINCE = time.time()
rows, cols = image.shape[0], image.shape[1]

num_point = 17

#creat original image mesh
src_cols = np.linspace(0, cols, num_point)
src_rows = np.linspace(0, rows, num_point)
src_rows, src_cols = np.meshgrid(src_rows, src_cols)
src = np.dstack([src_cols.flat, src_rows.flat])[0]

#create dst image mesh
dst_rows = np.array(src[:, 1])
dst_cols = np.array(src[:, 0])
#move a point by 9 pixel
dst_rows[6*num_point+9] += -9
dst = np.vstack([dst_cols, dst_rows]).T

#estimate transformation from the given src and dst mesh
tform = PiecewiseAffineTransform()
tform.estimate(src, dst)

out_rows = rows
out_cols = cols

#transform image with the estimated transformation
out = warp(image, tform, output_shape=(out_rows, out_cols))
#processing time
print(time.time()-SINCE)

#display image
fig, ax = plt.subplots(1, 2, figsize=(18, 9))
ax[0].imshow(out)
ax[0].plot(tform.inverse(src)[:, 0], tform.inverse(src)[:, 1], '.b')
ax[0].axis((0, out_cols, out_rows, 0))

ax[1].imshow(image)
ax[1].plot(src[:, 0], src[:, 1], '.b')
ax[1].axis((0, out_cols, out_rows, 0))

plt.show()
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/2/92796b3332a2fe74f867efc3aae2e5385a42b928.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/92796b3332a2fe74f867efc3aae2e5385a42b928" title="1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/2/92796b3332a2fe74f867efc3aae2e5385a42b928_2_517x240.jpeg" alt="1" data-base62-sha1="kTLQEa2ikwhDL6GsCW8Kv1vRSqs" width="517" height="240" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/2/92796b3332a2fe74f867efc3aae2e5385a42b928_2_517x240.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/9/2/92796b3332a2fe74f867efc3aae2e5385a42b928_2_775x360.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/2/92796b3332a2fe74f867efc3aae2e5385a42b928_2_1034x480.jpeg 2x" data-dominant-color="997F70"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">1</span><span class="informations">2330×1083 204 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The result is shown above, the red rectangle highlights the selected point.  However, it takes <strong>0.2s</strong> to finish the transformation. For a large dataset (~25, 000 images), it will take an extra ~1.4h to process images per epoch, which heavily increase the whole training time. Thus, I wonder if it can be accelerated if it can be processed on GPU.  But I can only find some random affine transformations in <code>torchvision.transforms.RandomAffine</code>. Is there any similar function that can transform an image with the given src and dst mesh in pytorch?</p>
<p>Thank you!</p>

username: ptrblck
cooked: <p>Maybe some transformations from <code>kornia.geometry.warp</code> could work.<br>
CC <a class="mention" href="/u/edgarriba">@edgarriba</a> as a <code>kornia</code> core dev.</p>

username: RuoyuGuo
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="177153">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p><code>kornia.geometry.warp</code></p>
</blockquote>
</aside>
<p>Thank you ptrblck! I will try it first!</p>

username: edgarriba
cooked: <p>You might need something like <a href="https://kornia.readthedocs.io/en/latest/geometry.transform.html#kornia.geometry.transform.warp_points_tps" class="inline-onebox" rel="noopener nofollow ugc">kornia.geometry.transform - Kornia</a></p>
<p>In any case we’ll check in detail that paper</p>