Spaces:
Build error
Build error
File size: 2,007 Bytes
ad1ac8f 72a1628 ad1ac8f 4e64649 ad1ac8f 4e64649 ad1ac8f a37eb28 ad1ac8f 79c6687 ad1ac8f 79c6687 ad1ac8f 72a1628 e4fb230 ad1ac8f 72a1628 ad1ac8f 79c6687 ad1ac8f 79c6687 ad1ac8f 72a1628 ad1ac8f e4fb230 |
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 |
# 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)
|