import torch


def im2col(img, win_len, stride=1):
	"""
	INPUTS:
	 - img: a b*c*h*w feature tensor.
	 - win_len: each pixel compares with its neighbors within a
				(win_len*2+1) * (win_len*2+1) window.
	OUTPUT:
	 - result: a b*c*(h*w)*(win_len*2+1)^2 tensor, unfolded neighbors for each pixel
	"""
	b,c,_,_ = img.size()
	# b * (c*w*w) * win_num
	unfold_img = torch.nn.functional.unfold(img, win_len, padding=0, stride=stride)
	unfold_img = unfold_img.view(b,c,win_len*win_len,-1)
	unfold_img = unfold_img.permute(0,1,3,2)
	return unfold_img