Elastic transformation/deformation of an image in Torch

elastic.lua

require 'image'

function ElasticTransform(img, alpha, sigma)
	--[[
	Parameters
	----------
	img: Tensor of size KxHxW
		Image on which elastic transformation have to be applied
	alpha: number
		Intensity of the transformation
	sigma: number
		Sigma for smoothing the transformation. Larger alphas require larger sigmas

	Returns
	-------
	img: Tensor of size KxHxW
		Image with elastic deformations appiled
	--]]

	H = img:size(2)
	W = img:size(3)
	filterSize = math.max(5,math.ceil(3*sigma))

	flow = torch.rand(2, H, W)*2 - 1 
	kernel = image.gaussian(filterSize, sigma, 1, true)
	flow = image.convolve(flow, kernel, 'same')*alpha

	img = image.warp(img, flow)
	return img
end