dataroots/SofaStyler

1# Import libraries2 3import cv24from tensorflow import keras5import numpy as np6from PIL import Image7import segmentation_models as sm8 9sm.set_framework("tf.keras")10 11# Load segmentation model12BACKBONE = "resnet50"13preprocess_input = sm.get_preprocessing(BACKBONE)14model = keras.models.load_model("Segmentation/model_final.h5", compile=False)15 16 17def get_mask(image: Image) -> Image:18    """19    This function generates a mask of the image that highlights all the sofas20    in the image. This uses a pre-trained Unet model with a resnet50 backbone.21    Remark: The model was trained on 640by640 images and it is therefore best22    that the image has the same size.23 24    Parameters:25            image = original image26    Return:27            mask  = corresponding maks of the image28    """29    test_img = np.array(image)30    test_img = cv2.resize(test_img, (640, 640))31    test_img = cv2.cvtColor(test_img, cv2.COLOR_RGB2BGR)32    test_img = np.expand_dims(test_img, axis=0)33 34    prediction = model.predict(preprocess_input(np.array(test_img))).round()35    mask = Image.fromarray(prediction[..., 0].squeeze() * 255).convert("L")36    return mask37 38 39def replace_sofa(image: Image, mask: Image, styled_sofa: Image) -> Image:40    """41    This function replaces the original sofa in the image by the new styled42    sofa according to the mask.43    Remark: All images should have the same size.44    Input:45        image       = Original image46        mask        = Generated masks highlighting the sofas in the image47        styled_sofa = Styled image48    Return:49        new_image   = Image containing the styled sofa50    """51    image, mask, styled_sofa = np.array(image), np.array(mask), np.array(styled_sofa)52 53    _, mask = cv2.threshold(mask, 10, 255, cv2.THRESH_BINARY)54    mask_inv = cv2.bitwise_not(mask)55    image_bg = cv2.bitwise_and(image, image, mask=mask_inv)56    sofa_fg = cv2.bitwise_and(styled_sofa, styled_sofa, mask=mask)57    new_image = cv2.add(image_bg, sofa_fg)58    return Image.fromarray(new_image)59