├── images ├── wandb_run.png ├── model_unet_like.png ├── model_partial_conv.png └── autoencoder_decoder.png ├── README.md └── utils └── pconv_layer.py /images/wandb_run.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ayulockin/deepimageinpainting/HEAD/images/wandb_run.png -------------------------------------------------------------------------------- /images/model_unet_like.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ayulockin/deepimageinpainting/HEAD/images/model_unet_like.png -------------------------------------------------------------------------------- /images/model_partial_conv.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ayulockin/deepimageinpainting/HEAD/images/model_partial_conv.png -------------------------------------------------------------------------------- /images/autoencoder_decoder.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ayulockin/deepimageinpainting/HEAD/images/autoencoder_decoder.png -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Introduction to image inpainting with deep learning 2 | 3 | By **Ayush Thakur** and [**Sayak Paul**](http://github.com/sayakpaul) 4 | 5 | This repository contains the supplementary notebook for the [Introduction to image impainting with deep learning](https://www.wandb.com/articles/introduction-to-image-inpainting-with-deep-learning) (from [Weights and Biases](https://www.wandb.com/)) article. 6 | 7 | ![](https://github.com/ayulockin/deepimageinpainting/blob/master/images/wandb_run.png?raw=true) 8 | 9 |
Weights and Biases run page
10 | 11 | 12 | ## Architectures 13 | 14 | #### UNET like network with standard convolutions 15 | 16 | ![UNET like with standard conv](https://github.com/ayulockin/deepimageimpainting/blob/master/images/model_unet_like.png) 17 | 18 | #### UNET like network with _partial convolutions_ 19 | 20 | ![UNET like with standard conv](https://github.com/ayulockin/deepimageimpainting/blob/master/images/model_partial_conv.png) 21 | -------------------------------------------------------------------------------- /utils/pconv_layer.py: -------------------------------------------------------------------------------- 1 | ## Reference: https://github.com/MathiasGruber/PConv-Keras/blob/master/libs/pconv_layer.py 2 | from tensorflow.keras import backend as K 3 | from tensorflow.keras.layers import InputSpec 4 | from tensorflow.keras.layers import Conv2D 5 | 6 | 7 | class PConv2D(Conv2D): 8 | def __init__(self, *args, n_channels=3, mono=False, **kwargs): 9 | super().__init__(*args, **kwargs) 10 | self.input_spec = [InputSpec(ndim=4), InputSpec(ndim=4)] 11 | 12 | def build(self, input_shape): 13 | """Adapted from original _Conv() layer of Keras 14 | param input_shape: list of dimensions for [img, mask] 15 | """ 16 | 17 | if self.data_format == 'channels_first': 18 | channel_axis = 1 19 | else: 20 | channel_axis = -1 21 | 22 | if input_shape[0][channel_axis] is None: 23 | raise ValueError('The channel dimension of the inputs should be defined. Found `None`.') 24 | 25 | self.input_dim = input_shape[0][channel_axis] 26 | 27 | # Image kernel 28 | kernel_shape = self.kernel_size + (self.input_dim, self.filters) 29 | self.kernel = self.add_weight(shape=kernel_shape, 30 | initializer=self.kernel_initializer, 31 | name='img_kernel', 32 | regularizer=self.kernel_regularizer, 33 | constraint=self.kernel_constraint) 34 | # Mask kernel 35 | self.kernel_mask = K.ones(shape=self.kernel_size + (self.input_dim, self.filters)) 36 | 37 | # Calculate padding size to achieve zero-padding 38 | self.pconv_padding = ( 39 | (int((self.kernel_size[0]-1)/2), int((self.kernel_size[0]-1)/2)), 40 | (int((self.kernel_size[0]-1)/2), int((self.kernel_size[0]-1)/2)), 41 | ) 42 | 43 | # Window size - used for normalization 44 | self.window_size = self.kernel_size[0] * self.kernel_size[1] 45 | 46 | if self.use_bias: 47 | self.bias = self.add_weight(shape=(self.filters,), 48 | initializer=self.bias_initializer, 49 | name='bias', 50 | regularizer=self.bias_regularizer, 51 | constraint=self.bias_constraint) 52 | else: 53 | self.bias = None 54 | self.built = True 55 | 56 | def call(self, inputs, mask=None): 57 | ''' 58 | We will be using the Keras conv2d method, and essentially we have 59 | to do here is multiply the mask with the input X, before we apply the 60 | convolutions. For the mask itself, we apply convolutions with all weights 61 | set to 1. 62 | Subsequently, we clip mask values to between 0 and 1 63 | ''' 64 | 65 | # Both image and mask must be supplied 66 | if type(inputs) is not list or len(inputs) != 2: 67 | raise Exception('PartialConvolution2D must be called on a list of two tensors [img, mask]. Instead got: ' + str(inputs)) 68 | 69 | # Padding done explicitly so that padding becomes part of the masked partial convolution 70 | images = K.spatial_2d_padding(inputs[0], self.pconv_padding, self.data_format) 71 | masks = K.spatial_2d_padding(inputs[1], self.pconv_padding, self.data_format) 72 | 73 | # Apply convolutions to mask 74 | mask_output = K.conv2d( 75 | masks, self.kernel_mask, 76 | strides=self.strides, 77 | padding='valid', 78 | data_format=self.data_format, 79 | dilation_rate=self.dilation_rate 80 | ) 81 | 82 | # Apply convolutions to image 83 | img_output = K.conv2d( 84 | (images*masks), self.kernel, 85 | strides=self.strides, 86 | padding='valid', 87 | data_format=self.data_format, 88 | dilation_rate=self.dilation_rate 89 | ) 90 | 91 | # Calculate the mask ratio on each pixel in the output mask 92 | mask_ratio = self.window_size / (mask_output + 1e-8) 93 | 94 | # Clip output to be between 0 and 1 95 | mask_output = K.clip(mask_output, 0, 1) 96 | 97 | # Remove ratio values where there are holes 98 | mask_ratio = mask_ratio * mask_output 99 | 100 | # Normalize iamge output 101 | img_output = img_output * mask_ratio 102 | 103 | # Apply bias only to the image (if chosen to do so) 104 | if self.use_bias: 105 | img_output = K.bias_add( 106 | img_output, 107 | self.bias, 108 | data_format=self.data_format) 109 | 110 | # Apply activations on the image 111 | if self.activation is not None: 112 | img_output = self.activation(img_output) 113 | 114 | return [img_output, mask_output] 115 | 116 | def compute_output_shape(self, input_shape): 117 | if self.data_format == 'channels_last': 118 | space = input_shape[0][1:-1] 119 | new_space = [] 120 | for i in range(len(space)): 121 | new_dim = conv_output_length( 122 | space[i], 123 | self.kernel_size[i], 124 | padding='same', 125 | stride=self.strides[i], 126 | dilation=self.dilation_rate[i]) 127 | new_space.append(new_dim) 128 | new_shape = (input_shape[0][0],) + tuple(new_space) + (self.filters,) 129 | return [new_shape, new_shape] 130 | if self.data_format == 'channels_first': 131 | space = input_shape[2:] 132 | new_space = [] 133 | for i in range(len(space)): 134 | new_dim = conv_output_length( 135 | space[i], 136 | self.kernel_size[i], 137 | padding='same', 138 | stride=self.strides[i], 139 | dilation=self.dilation_rate[i]) 140 | new_space.append(new_dim) 141 | new_shape = (input_shape[0], self.filters) + tuple(new_space) 142 | return [new_shape, new_shape] 143 | 144 | ## Reference: https://github.com/keras-team/keras/blob/7a39b6c62d43c25472b2c2476bd2a8983ae4f682/keras/utils/conv_utils.py#L85 145 | def conv_output_length(input_length, filter_size, 146 | padding, stride, dilation=1): 147 | """Determines output length of a convolution given input length. 148 | # Arguments 149 | input_length: integer. 150 | filter_size: integer. 151 | padding: one of `"same"`, `"valid"`, `"full"`. 152 | stride: integer. 153 | dilation: dilation rate, integer. 154 | # Returns 155 | The output length (integer). 156 | """ 157 | if input_length is None: 158 | return None 159 | assert padding in {'same', 'valid', 'full', 'causal'} 160 | dilated_filter_size = (filter_size - 1) * dilation + 1 161 | if padding == 'same': 162 | output_length = input_length 163 | elif padding == 'valid': 164 | output_length = input_length - dilated_filter_size + 1 165 | elif padding == 'causal': 166 | output_length = input_length 167 | elif padding == 'full': 168 | output_length = input_length + dilated_filter_size - 1 169 | return (output_length + stride - 1) // stride 170 | --------------------------------------------------------------------------------