init space

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+.idea/
+__pycache__/

app.py

@@ -0,0 +1,41 @@
+import gradio as gr
+from bald_face import BaldFace
+
+bald = BaldFace()
+
+
+def predict(image):
+    return bald.make(image)
+
+
+footer = r"""
+<center>
+<b>
+Demo for <a href=''>Lightweight OpenPose</a>
+</b>
+</center>
+"""
+
+with gr.Blocks(title="") as app:
+    gr.HTML("<center><h1></h1></center>")
+    gr.HTML("<center><h3></h3></center>")
+    with gr.Row(equal_height=False):
+        with gr.Column():
+            input_img = gr.Image(type="numpy", label="Input image")
+            run_btn = gr.Button(variant="primary")
+        with gr.Column():
+            output_img = gr.Image(type="pil", label="Output image")
+            gr.ClearButton(components=[input_img, output_img], variant="stop")
+
+    run_btn.click(predict, [input_img], [output_img])
+
+    with gr.Row():
+        blobs = [[f"examples/{x:02d}.jpg"] for x in range(1, 5)]
+        examples = gr.Dataset(components=[input_img], samples=blobs)
+        examples.click(lambda x: x[0], [examples], [input_img])
+
+    with gr.Row():
+        gr.HTML(footer)
+
+app.launch(share=False, debug=True, show_error=True)
+app.queue()

bald_face.py

@@ -0,0 +1,109 @@
+#############################################################################
+#
+#   Source from:
+#   https://github.com/leonelhs/baldgan
+#   Forked from:
+#   https://github.com/david-svitov/baldgan
+#   Reimplemented by: Leonel Hernández
+#
+##############################################################################
+import os
+
+import PIL.Image
+import cv2
+import numpy as np
+from retinaface import RetinaFace
+from skimage import transform as trans
+from huggingface_hub import hf_hub_download
+
+from baldgan.model import buildModel
+
+BALDGAN_REPO_ID = "leonelhs/baldgan"
+
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"  # see issue #152
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+gpu_id = -1
+
+image_size = [256, 256]
+
+src = np.array([
+    [30.2946, 51.6963],
+    [65.5318, 51.5014],
+    [48.0252, 71.7366],
+    [33.5493, 92.3655],
+    [62.7299, 92.2041]], dtype=np.float32)
+
+src[:, 0] += 8.0
+src[:, 0] += 15.0
+src[:, 1] += 30.0
+src /= 112
+src *= 200
+
+
+def list2array(values):
+    return np.array(list(values))
+
+
+def align_face(img):
+    faces = RetinaFace.detect_faces(img)
+    bounding_boxes = np.array([list2array(faces[face]['facial_area']) for face in faces])
+    points = np.array([list2array(faces[face]['landmarks'].values()) for face in faces])
+    white_image = np.ones(img.shape, dtype=np.uint8) * 255
+
+    result_faces = []
+    result_masks = []
+    result_matrix = []
+
+    if bounding_boxes.shape[0] > 0:
+        det = bounding_boxes[:, 0:4]
+        for i in range(det.shape[0]):
+            _det = det[i]
+            dst = points[i]
+
+            tform = trans.SimilarityTransform()
+            tform.estimate(dst, src)
+            M = tform.params[0:2, :]
+            warped = cv2.warpAffine(img, M, (image_size[1], image_size[0]), borderValue=0.0)
+            mask = cv2.warpAffine(white_image, M, (image_size[1], image_size[0]), borderValue=0.0)
+
+            result_faces.append(warped)
+            result_masks.append(mask)
+            result_matrix.append(tform.params[0:3, :])
+
+    return result_faces, result_masks, result_matrix
+
+
+def put_face_back(img, faces, masks, result_matrix):
+    for i in range(len(faces)):
+        M = np.linalg.inv(result_matrix[i])[0:2]
+        warped = cv2.warpAffine(faces[i], M, (img.shape[1], img.shape[0]), borderValue=0.0)
+        mask = cv2.warpAffine(masks[i], M, (img.shape[1], img.shape[0]), borderValue=0.0)
+        mask = mask // 255
+        img = img * (1 - mask)
+        img = img.astype(np.uint8)
+        img += warped * mask
+    return img
+
+
+class BaldFace:
+
+    def __init__(self):
+        self.model = buildModel()
+        model_path = hf_hub_download(repo_id=BALDGAN_REPO_ID, filename='model_G_5_170.hdf5')
+        self.model.load_weights(model_path)
+
+    def make(self, image):
+        # image = np.array(image)
+        faces, masks, matrix = align_face(image)
+        result_faces = []
+
+        for face in faces:
+            input_face = np.expand_dims(face, axis=0)
+            input_face = input_face / 127.5 - 1.
+            result = self.model.predict(input_face)[0]
+            result = ((result + 1.) * 127.5)
+            result = result.astype(np.uint8)
+            result_faces.append(result)
+
+        img_result = put_face_back(image, result_faces, masks, matrix)
+        return PIL.Image.fromarray(img_result)

baldgan/keras_contrib/layers/normalization/instancenormalization.py

@@ -0,0 +1,147 @@
+from keras.layers import Layer, InputSpec
+from keras import initializers, regularizers, constraints
+from keras import backend as K
+
+
+class InstanceNormalization(Layer):
+    """Instance normalization layer.
+
+    Normalize the activations of the previous layer at each step,
+    i.e. applies a transformation that maintains the mean activation
+    close to 0 and the activation standard deviation close to 1.
+
+    # Arguments
+        axis: Integer, the axis that should be normalized
+            (typically the features axis).
+            For instance, after a `Conv2D` layer with
+            `data_format="channels_first"`,
+            set `axis=1` in `InstanceNormalization`.
+            Setting `axis=None` will normalize all values in each
+            instance of the batch.
+            Axis 0 is the batch dimension. `axis` cannot be set to 0 to avoid errors.
+        epsilon: Small float added to variance to avoid dividing by zero.
+        center: If True, add offset of `beta` to normalized tensor.
+            If False, `beta` is ignored.
+        scale: If True, multiply by `gamma`.
+            If False, `gamma` is not used.
+            When the next layer is linear (also e.g. `nn.relu`),
+            this can be disabled since the scaling
+            will be done by the next layer.
+        beta_initializer: Initializer for the beta weight.
+        gamma_initializer: Initializer for the gamma weight.
+        beta_regularizer: Optional regularizer for the beta weight.
+        gamma_regularizer: Optional regularizer for the gamma weight.
+        beta_constraint: Optional constraint for the beta weight.
+        gamma_constraint: Optional constraint for the gamma weight.
+
+    # Input shape
+        Arbitrary. Use the keyword argument `input_shape`
+        (tuple of integers, does not include the samples axis)
+        when using this layer as the first layer in a Sequential model.
+
+    # Output shape
+        Same shape as input.
+
+    # References
+        - [Layer Normalization](https://arxiv.org/abs/1607.06450)
+        - [Instance Normalization: The Missing Ingredient for Fast Stylization](
+        https://arxiv.org/abs/1607.08022)
+    """
+    def __init__(self,
+                 axis=None,
+                 epsilon=1e-3,
+                 center=True,
+                 scale=True,
+                 beta_initializer='zeros',
+                 gamma_initializer='ones',
+                 beta_regularizer=None,
+                 gamma_regularizer=None,
+                 beta_constraint=None,
+                 gamma_constraint=None,
+                 **kwargs):
+        super(InstanceNormalization, self).__init__(**kwargs)
+        self.supports_masking = True
+        self.axis = axis
+        self.epsilon = epsilon
+        self.center = center
+        self.scale = scale
+        self.beta_initializer = initializers.get(beta_initializer)
+        self.gamma_initializer = initializers.get(gamma_initializer)
+        self.beta_regularizer = regularizers.get(beta_regularizer)
+        self.gamma_regularizer = regularizers.get(gamma_regularizer)
+        self.beta_constraint = constraints.get(beta_constraint)
+        self.gamma_constraint = constraints.get(gamma_constraint)
+
+    def build(self, input_shape):
+        ndim = len(input_shape)
+        if self.axis == 0:
+            raise ValueError('Axis cannot be zero')
+
+        if (self.axis is not None) and (ndim == 2):
+            raise ValueError('Cannot specify axis for rank 1 tensor')
+
+        self.input_spec = InputSpec(ndim=ndim)
+
+        if self.axis is None:
+            shape = (1,)
+        else:
+            shape = (input_shape[self.axis],)
+
+        if self.scale:
+            self.gamma = self.add_weight(shape=shape,
+                                         name='gamma',
+                                         initializer=self.gamma_initializer,
+                                         regularizer=self.gamma_regularizer,
+                                         constraint=self.gamma_constraint)
+        else:
+            self.gamma = None
+        if self.center:
+            self.beta = self.add_weight(shape=shape,
+                                        name='beta',
+                                        initializer=self.beta_initializer,
+                                        regularizer=self.beta_regularizer,
+                                        constraint=self.beta_constraint)
+        else:
+            self.beta = None
+        self.built = True
+
+    def call(self, inputs, training=None):
+        input_shape = K.int_shape(inputs)
+        reduction_axes = list(range(0, len(input_shape)))
+
+        if self.axis is not None:
+            del reduction_axes[self.axis]
+
+        del reduction_axes[0]
+
+        mean = K.mean(inputs, reduction_axes, keepdims=True)
+        stddev = K.std(inputs, reduction_axes, keepdims=True) + self.epsilon
+        normed = (inputs - mean) / stddev
+
+        broadcast_shape = [1] * len(input_shape)
+        if self.axis is not None:
+            broadcast_shape[self.axis] = input_shape[self.axis]
+
+        if self.scale:
+            broadcast_gamma = K.reshape(self.gamma, broadcast_shape)
+            normed = normed * broadcast_gamma
+        if self.center:
+            broadcast_beta = K.reshape(self.beta, broadcast_shape)
+            normed = normed + broadcast_beta
+        return normed
+
+    def get_config(self):
+        config = {
+            'axis': self.axis,
+            'epsilon': self.epsilon,
+            'center': self.center,
+            'scale': self.scale,
+            'beta_initializer': initializers.serialize(self.beta_initializer),
+            'gamma_initializer': initializers.serialize(self.gamma_initializer),
+            'beta_regularizer': regularizers.serialize(self.beta_regularizer),
+            'gamma_regularizer': regularizers.serialize(self.gamma_regularizer),
+            'beta_constraint': constraints.serialize(self.beta_constraint),
+            'gamma_constraint': constraints.serialize(self.gamma_constraint)
+        }
+        base_config = super(InstanceNormalization, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))

baldgan/model.py

@@ -0,0 +1,97 @@
+import keras.backend as K
+from keras import Model
+from keras.layers import GlobalAveragePooling2D, multiply, LeakyReLU, Permute
+from keras.layers import Input, Dense, Reshape, Dropout, Concatenate
+from keras.src.layers import UpSampling2D, Conv2D
+from .keras_contrib.layers.normalization.instancenormalization import InstanceNormalization
+
+
+def squeeze_excite_block(input, ratio=4):
+    """ Create a channel-wise squeeze-excite block
+    Args:
+        input: input tensor
+        filters: number of output filters
+    Returns: a keras tensor
+    References
+    -   [Squeeze and Excitation Networks](https://arxiv.org/abs/1709.01507)
+    :param input:
+    :param ratio:
+    """
+    init = input
+    channel_axis = 1 if K.image_data_format() == "channels_first" else -1
+    filters = init.shape[channel_axis]
+    se_shape = (1, 1, filters)
+
+    se = GlobalAveragePooling2D()(init)
+    se = Reshape(se_shape)(se)
+    se = Dense(filters // ratio, activation='relu', kernel_initializer='he_normal', use_bias=False)(se)
+    se = Dense(filters, activation='sigmoid', kernel_initializer='he_normal', use_bias=False)(se)
+
+    if K.image_data_format() == 'channels_first':
+        se = Permute((3, 1, 2))(se)
+
+    x = multiply([init, se])
+    return x
+
+
+def conv2d(layer_input, filters, f_size=4, bn=True, se=False):
+    """Layers used during down sampling"""
+    d = Conv2D(filters, kernel_size=f_size, strides=2, padding='same')(layer_input)
+    d = LeakyReLU(alpha=0.2)(d)
+    if bn:
+        d = InstanceNormalization()(d)
+    if se:
+        d = squeeze_excite_block(d)
+    return d
+
+
+def deconv2d(layer_input, skip_input, filters, f_size=4, dropout_rate=0):
+    """Layers used during up sampling"""
+    u = UpSampling2D(size=2)(layer_input)
+    u = Conv2D(filters, kernel_size=f_size, strides=1, padding='same', activation='relu')(u)
+    if dropout_rate:
+        u = Dropout(dropout_rate)(u)
+    u = InstanceNormalization()(u)
+    u = Concatenate()([u, skip_input])
+    return u
+
+
+def atrous(layer_input, filters, f_size=4, bn=True):
+    a_list = []
+    for rate in [2, 4, 8]:
+        a = Conv2D(filters, f_size, dilation_rate=rate, padding='same')(layer_input)
+        a_list.append(a)
+    a = Concatenate()(a_list)
+    a = LeakyReLU(alpha=0.2)(a)
+    if bn:
+        a = InstanceNormalization()(a)
+    return a
+
+
+def buildModel():
+    K.set_learning_phase(0)
+
+    # Image input
+    d0 = Input(shape=(256, 256, 3))
+
+    gf = 64
+    # Down sampling
+    d1 = conv2d(d0, gf, bn=False, se=True)
+    d2 = conv2d(d1, gf * 2, se=True)
+    d3 = conv2d(d2, gf * 4, se=True)
+    d4 = conv2d(d3, gf * 8)
+    d5 = conv2d(d4, gf * 8)
+
+    a1 = atrous(d5, gf * 8)
+
+    # Up sampling
+    u3 = deconv2d(a1, d4, gf * 8)
+    u4 = deconv2d(u3, d3, gf * 4)
+    u5 = deconv2d(u4, d2, gf * 2)
+    u6 = deconv2d(u5, d1, gf)
+
+    u7 = UpSampling2D(size=2)(u6)
+
+    output_img = Conv2D(3, kernel_size=4, strides=1, padding='same', activation='tanh')(u7)
+
+    return Model(d0, output_img)

requirements.txt

@@ -0,0 +1,6 @@
+retina-face
+keras~=2.15.0
+pillow~=10.1.0
+opencv-python~=4.8.1.78
+numpy~=1.26.2
+scikit-image~=0.22.0