Spaces:
Build error
Build error
| # Import libraries | |
| import cv2 | |
| from tensorflow import keras | |
| import numpy as np | |
| from PIL import Image | |
| import segmentation_models as sm | |
| sm.set_framework("tf.keras") | |
| # Load segmentation model | |
| BACKBONE = "resnet50" | |
| preprocess_input = sm.get_preprocessing(BACKBONE) | |
| model = keras.models.load_model("Segmentation/model_final.h5", compile=False) | |
| def get_mask(image: Image) -> Image: | |
| """ | |
| This function generates a mask of the image that highlights all the sofas | |
| in the image. This uses a pre-trained Unet model with a resnet50 backbone. | |
| Remark: The model was trained on 640by640 images and it is therefore best | |
| that the image has the same size. | |
| Parameters: | |
| image = original image | |
| Return: | |
| mask = corresponding maks of the image | |
| """ | |
| test_img = np.array(image) | |
| test_img = cv2.resize(test_img, (640, 640)) | |
| test_img = cv2.cvtColor(test_img, cv2.COLOR_RGB2BGR) | |
| test_img = np.expand_dims(test_img, axis=0) | |
| prediction = model.predict(preprocess_input(np.array(test_img))).round() | |
| mask = Image.fromarray(prediction[..., 0].squeeze() * 255).convert("L") | |
| return mask | |
| def replace_sofa(image: Image, mask: Image, styled_sofa: Image) -> Image: | |
| """ | |
| This function replaces the original sofa in the image by the new styled | |
| sofa according to the mask. | |
| Remark: All images should have the same size. | |
| Input: | |
| image = Original image | |
| mask = Generated masks highlighting the sofas in the image | |
| styled_sofa = Styled image | |
| Return: | |
| new_image = Image containing the styled sofa | |
| """ | |
| image, mask, styled_sofa = np.array(image), np.array(mask), np.array(styled_sofa) | |
| _, mask = cv2.threshold(mask, 10, 255, cv2.THRESH_BINARY) | |
| mask_inv = cv2.bitwise_not(mask) | |
| image_bg = cv2.bitwise_and(image, image, mask=mask_inv) | |
| sofa_fg = cv2.bitwise_and(styled_sofa, styled_sofa, mask=mask) | |
| new_image = cv2.add(image_bg, sofa_fg) | |
| return Image.fromarray(new_image) | |