Spaces:

zama-fhe
/

encrypted_image_filtering

Running

App Files Files Community

encrypted_image_filtering / filters.py

romanbredehoft-zama

Update Concrete ML version and re-generate development files

a13c444 9 months ago

raw history blame

No virus

8.75 kB

	"Filter definitions, with pre-processing, post-processing and compilation methods."

	import numpy as np
	import torch
	from torch import nn
	from common import AVAILABLE_FILTERS, INPUT_SHAPE

	from concrete.fhe.compilation.compiler import Compiler
	from concrete.ml.common.utils import generate_proxy_function
	from concrete.ml.torch.numpy_module import NumpyModule


	class TorchIdentity(nn.Module):
	"""Torch identity model."""

	def forward(self, x):
	"""Identity forward pass.

	Args:
	x (torch.Tensor): The input image.

	Returns:
	x (torch.Tensor): The input image.
	"""
	return x


	class TorchInverted(nn.Module):
	"""Torch inverted model."""

	def forward(self, x):
	"""Forward pass for inverting an image's colors.

	Args:
	x (torch.Tensor): The input image.

	Returns:
	torch.Tensor: The (color) inverted image.
	"""
	return 255 - x


	class TorchRotate(nn.Module):
	"""Torch rotated model."""

	def forward(self, x):
	"""Forward pass for rotating an image.

	Args:
	x (torch.Tensor): The input image.

	Returns:
	torch.Tensor: The rotated image.
	"""
	return x.transpose(0, 1)


	class TorchConv(nn.Module):
	"""Torch model with a single convolution operator."""

	def __init__(self, kernel, n_in_channels=3, n_out_channels=3, groups=1, threshold=None):
	"""Initialize the filter.

	Args:
	kernel (np.ndarray): The convolution kernel to consider.
	"""
	super().__init__()
	self.kernel = torch.tensor(kernel, dtype=torch.int64)
	self.n_out_channels = n_out_channels
	self.n_in_channels = n_in_channels
	self.groups = groups
	self.threshold = threshold

	def forward(self, x):
	"""Forward pass with a single convolution using a 1D or 2D kernel.

	Args:
	x (torch.Tensor): The input image.

	Returns:
	torch.Tensor: The filtered image.
	"""
	# Define the convolution parameters
	stride = 1
	kernel_shape = self.kernel.shape

	# Ensure the kernel has a proper shape
	# If the kernel has a 1D shape, a (1, 1) kernel is used for each in_channels
	if len(kernel_shape) == 1:
	self.kernel = self.kernel.repeat(self.n_out_channels)
	kernel = self.kernel.reshape(
	self.n_out_channels,
	self.n_in_channels // self.groups,
	1,
	1,
	)

	# Else, if the kernel has a 2D shape, a single (Kw, Kh) kernel is used on all in_channels
	elif len(kernel_shape) == 2:
	kernel = self.kernel.expand(
	self.n_out_channels,
	self.n_in_channels // self.groups,
	kernel_shape[0],
	kernel_shape[1],
	)


	else:
	raise ValueError(
	"Wrong kernel shape, only 1D or 2D kernels are accepted. Got kernel of shape "
	f"{kernel_shape}"
	)

	# Reshape the image. This is done because Torch convolutions and Numpy arrays (for PIL
	# display) don't follow the same shape conventions. More precisely, x is of shape
	# (Width, Height, Channels) while the conv2d operator requires an input of shape
	# (Batch, Channels, Height, Width)
	x = x.transpose(2, 0).unsqueeze(axis=0)

	# Apply the convolution
	x = nn.functional.conv2d(x, kernel, stride=stride, groups=self.groups)

	# Reshape the output back to the original shape (Width, Height, Channels)
	x = x.transpose(1, 3).reshape((x.shape[2], x.shape[3], self.n_out_channels))

	# Subtract a given threshold if given
	if self.threshold is not None:
	x -= self.threshold

	return x


	class Filter:
	"""Filter class used in the app."""

	def __init__(self, filter_name):
	"""Initializing the filter class using a given filter.

	Most filters can be found at https://en.wikipedia.org/wiki/Kernel_(image_processing).

	Args:
	filter_name (str): The filter to consider.
	"""

	assert filter_name in AVAILABLE_FILTERS, (
	f"Unsupported image filter or transformation. Expected one of {*AVAILABLE_FILTERS,}, "
	f"but got {filter_name}",
	)

	# Define attributes associated to the filter
	self.filter_name = filter_name
	self.onnx_model = None
	self.fhe_circuit = None
	self.divide = None

	# Instantiate the torch module associated to the given filter name
	if filter_name == "identity":
	self.torch_model = TorchIdentity()

	elif filter_name == "inverted":
	self.torch_model = TorchInverted()

	elif filter_name == "rotate":
	self.torch_model = TorchRotate()

	elif filter_name == "black and white":
	# Define the grayscale weights (RGB order)
	# These weights were used in PAL and NTSC video systems and can be found at
	# https://en.wikipedia.org/wiki/Grayscale
	# There are initially supposed to be float weights (0.299, 0.587, 0.114), with
	# 0.299 + 0.587 + 0.114 = 1
	# However, since FHE computations require weights to be integers, we first multiply
	# these by a factor of 1000. The output image's values are then divided by 1000 in
	# post-processing in order to retrieve the correct result
	kernel = [299, 587, 114]

	self.torch_model = TorchConv(kernel)

	# Define the value used when for dividing the output values in post-processing
	self.divide = 1000


	elif filter_name == "blur":
	kernel = np.ones((3, 3))

	self.torch_model = TorchConv(kernel, groups=3)

	# Define the value used when for dividing the output values in post-processing
	self.divide = 9

	elif filter_name == "sharpen":
	kernel = [
	[0, -1, 0],
	[-1, 5, -1],
	[0, -1, 0],
	]

	self.torch_model = TorchConv(kernel, groups=3)

	elif filter_name == "ridge detection":
	kernel = [
	[-1, -1, -1],
	[-1, 9, -1],
	[-1, -1, -1],
	]

	# Additionally to the convolution operator, the filter will subtract a given threshold
	# value to the result in order to better display the ridges
	self.torch_model = TorchConv(kernel, threshold=900)


	def compile(self):
	"""Compile the filter on a representative inputset."""
	# Generate a random representative set of images used for compilation, following shape
	# PIL's shape RGB format for Numpy arrays (image_width, image_height, 3)
	# Additionally, this version's compiler only handles tuples of 1-batch array as inputset,
	# meaning we need to define the inputset as a Tuple[np.ndarray[shape=(H, W, 3)]]
	np.random.seed(42)
	inputset = tuple(
	np.random.randint(0, 256, size=(INPUT_SHAPE + (3, )), dtype=np.int64) for _ in range(100)
	)

	# Convert the Torch module to a Numpy module
	numpy_module = NumpyModule(
	self.torch_model,
	dummy_input=torch.from_numpy(inputset[0]),
	)

	# Get the proxy function and parameter mappings used for initializing the compiler
	# This is done in order to be able to provide any modules with arbitrary numbers of
	# encrypted arguments to Concrete Numpy's compiler
	numpy_filter_proxy, parameters_mapping = generate_proxy_function(
	numpy_module.numpy_forward,
	["inputs"]
	)

	# Compile the filter and retrieve its FHE circuit
	compiler = Compiler(
	numpy_filter_proxy,
	{parameters_mapping["inputs"]: "encrypted"},
	)
	self.fhe_circuit = compiler.compile(inputset)

	return self.fhe_circuit

	def post_processing(self, output_image):
	"""Apply post-processing to the encrypted output images.

	Args:
	input_image (np.ndarray): The decrypted image to post-process.

	Returns:
	input_image (np.ndarray): The post-processed image.
	"""
	# Divide all values if needed
	if self.divide is not None:
	output_image //= self.divide

	# Clip the image's values to proper RGB standards as filters don't handle such constraints
	output_image = output_image.clip(0, 255)

	return output_image