示例#1
0
def main(argv):
	parser = ArgumentParser(argv[0], description=__doc__)
	parser.add_argument('image',                    type=str)
	parser.add_argument('model',                    type=str)
	parser.add_argument('--init',             '-I', type=str, default=None)
	parser.add_argument('--index',            '-x', type=int, default=0,
		help='Determines which image is used when whole dataset is given instead of image.')
	parser.add_argument('--fill_region',      '-f', type=int, default=71)
	parser.add_argument('--outer_patch_size', '-p', type=int, default=19)
	parser.add_argument('--inner_patch_size', '-i', type=int, default=5)
	parser.add_argument('--stride',           '-s', type=int, default=3)
	parser.add_argument('--candidates',       '-C', type=int, default=5,
		help='The best initialization is taken out of this many initializations.')
	parser.add_argument('--num_epochs',       '-e', type=int, default=1000)
	parser.add_argument('--method',           '-m', type=str, default='SAMPLE', choices=['SAMPLE', 'MAP'])
	parser.add_argument('--step_width',       '-l', type=float, default=100.)
	parser.add_argument('--output',           '-o', type=str, default='results/inpainting/')
	parser.add_argument('--flip',             '-F', type=int, default=0,
		help='If > 0, assume horizontal symmetry. If > 1, assume vertical symmetry.')

	args = parser.parse_args(argv[1:])


	### DATA

	# load image
	if args.image.lower()[-4:] in ['.gif', '.png', '.jpg', 'jpeg']:
		image = imread(args.image)[None]
		vmin, vmax = 0, 255
	else:
		image = loadmat(args.image)['data'][[args.index]]
		vmin, vmax = image.min(), image.max()
	if image.ndim < 4:
		image = image[:, :, :, None]

	image = asarray(image, dtype=float)

	imwrite(os.path.join(args.output, 'original.png'),
		imformat(image[0, :, :, 0], vmin=vmin, vmax=vmax, symmetric=False))

	# remove center portion
	i_start = (image.shape[1] - args.fill_region) // 2
	j_start = (image.shape[2] - args.fill_region) // 2
	image[0,
		i_start:i_start + args.fill_region,
		j_start:j_start + args.fill_region, 0] = vmin + rand(args.fill_region, args.fill_region) * (vmax - vmin)

	imwrite(os.path.join(args.output, 'start.png'),
		imformat(image[0, :, :, 0], vmin=vmin, vmax=vmax, symmetric=False))


	### MODEL

	# load model
	model = Experiment(args.model)['model']
	model.verbosity = False

	# use different models for sampling and likelihoods because of SLSTM caching
	model_copy = deepcopy(model)

	# create mask indicating pixels to replace
	M = args.outer_patch_size
	N = args.inner_patch_size
	m = (M - N) // 2
	n = M - N - m
	patch_mask = zeros([M, M], dtype=bool)
	patch_mask[m:-n, m:-n] = True

	if args.init is None:
		candidates = []
		logliks = []

		for _ in range(args.candidates):
			# replace missing pixels by ancestral sampling
			patch = image[:,
				i_start - M:i_start + args.fill_region,
				j_start - M:j_start + args.fill_region + M]
			sample_mask = zeros([patch.shape[1], patch.shape[2]], dtype=bool)
			sample_mask[M:, M:-M] = True
			image[:,
				i_start - M:i_start + args.fill_region,
				j_start - M:j_start + args.fill_region + M] = model.sample(patch, mask=sample_mask,
					min_values=vmin, max_values=vmax)

			candidates.append(image.copy())
			logliks.append(model.loglikelihood(image).sum())

		image = candidates[argmax(logliks)]

		imwrite(os.path.join(args.output, 'fillin.0.png'),
			imformat(image[0, :, :, 0], vmin=vmin, vmax=vmax, symmetric=False))

		start_epoch = 0

	else:
		init = load(args.init)
		image = init['image']
		start_epoch = init['epoch']


	### INPAINTING

	try:
		for epoch in range(start_epoch, args.num_epochs):
			print epoch
			
			h_flipped = False
			if args.flip > 0 and rand() < .5:
				print 'Horizontal flip.'
				# flip image horizontally
				image = image[:, :, ::-1]
				j_start = image.shape[2] - j_start - args.fill_region
				h_flipped = True

			v_flipped = False
			if args.flip > 0 and rand() < .5:
				print 'Vertical flip.'
				# flip image vertically
				image = image[:, ::-1, :]
				i_start = image.shape[1] - i_start - args.fill_region
				v_flipped = True

			for i in range(i_start - m, i_start - m + args.fill_region - N + 1, args.stride):
				for j in range(j_start - m, j_start - m + args.fill_region - N + 1, args.stride):
					patch = image[:, i:i + M, j:j + M]

					if args.method == 'SAMPLE':
						# proposal
						patch_pr, logq_pr = model.sample(patch.copy(), mask=patch_mask,
							min_values=vmin, max_values=vmax, return_loglik=True)

						# conditional log-density
						logq = model_copy._logq(patch, patch_mask)

						# joint log-densities
						logp = model_copy.loglikelihood(patch).sum()
						logp_pr = model_copy.loglikelihood(patch_pr).sum()

						if rand() < exp(logp_pr - logp - logq_pr + logq):
							# accept proposal
							patch[:] = patch_pr

					else:
						# gradient step
						grad = model.gradient(patch)[1]
						patch[:, patch_mask] += grad[:, patch_mask] * args.step_width

			# flip back
			if h_flipped:
				image = image[:, :, ::-1]
				j_start = image.shape[2] - j_start - args.fill_region
			if v_flipped:
				image = image[:, ::-1, :]
				i_start = image.shape[1] - i_start - args.fill_region

			imwrite(os.path.join(args.output, 'fillin.{0}.png'.format(epoch + 1)),
				imformat(image[0, :, :, 0], vmin=vmin, vmax=vmax, symmetric=False))

	except KeyboardInterrupt:
		pass

	savez(os.path.join(args.output, 'final.npz'), image=image, epoch=epoch)

	return 0