File size: 8,745 Bytes
21c7197 7b32412 21c7197 a13c444 21c7197 7b32412 21c7197 7b32412 21c7197 7b32412 21c7197 1974e22 21c7197 7b32412 2a040cc 21c7197 3cf0931 21c7197 3cf0931 21c7197 3cf0931 21c7197 2a040cc 21c7197 d9d229c 21c7197 1974e22 21c7197 1974e22 3cf0931 1974e22 3cf0931 21c7197 4bee342 21c7197 4bee342 21c7197 7b32412 21c7197 4bee342 21c7197 4bee342 7b32412 3cf0931 21c7197 4bee342 7b32412 21c7197 4bee342 7b32412 21c7197 4bee342 7b32412 21c7197 4bee342 21c7197 3cf0931 21c7197 d9d229c 21c7197 3cf0931 21c7197 4bee342 3cf0931 21c7197 d9d229c 21c7197 3cf0931 21c7197 4bee342 3cf0931 21c7197 d9d229c 21c7197 4bee342 3cf0931 d9d229c 3cf0931 21c7197 2a040cc 02841d1 3cf0931 21c7197 1974e22 21c7197 2a040cc 21c7197 2a040cc 21c7197 2a040cc 21c7197 3cf0931 21c7197 3cf0931 21c7197 3cf0931 21c7197 3cf0931 21c7197 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 |
"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
|