more testing

app.py

@@ -98,8 +98,8 @@ def style_sofa(input_img: np.ndarray, style_img: np.ndarray):
     resized_style = resize_style(style_img)
     #resized_style.save('resized_style.jpg')
     # generate mask for image
-    # print('generating mask...')
-    # mask = get_mask(resized_img)
+    print('generating mask...')
+    mask = get_mask(resized_img)
     #mask.save('mask.jpg')
     # Created a styled sofa
     # print('Styling sofa...')
@@ -110,7 +110,7 @@ def style_sofa(input_img: np.ndarray, style_img: np.ndarray):
     # new_sofa = replace_sofa(resized_img,mask,styled_sofa)
     # new_sofa = new_sofa.crop(box)
     print('Finishing job', id)
-    return resized_img
+    return mask
 
 demo = gr.Interface(
     style_sofa,

segmentation.py

@@ -6,6 +6,7 @@ import numpy as np
 import matplotlib.pyplot as plt
 from PIL import Image
 import segmentation_models as sm
+sm.set_framework('tf.keras')
 
 
 def get_mask(image:Image) -> Image:
@@ -27,7 +28,6 @@ def get_mask(image:Image) -> Image:
     # in_classes = 1 if len(CLASSES) == 1 else (len(CLASSES) + 1)  # case for binary and multiclass segmentation
     # actvation = 'sigmoid' if n_classes == 1 else 'softmax'
     preprocess_input = sm.get_preprocessing(BACKBONE)
-    sm.set_framework('tf.keras')
     # LR=0.0001
 
     #create model architecture
@@ -47,9 +47,10 @@ def get_mask(image:Image) -> Image:
     # #load model
     # model.load_weights(model_path) 
     model = keras.models.load_model('model_final.h5', compile=False)
-
+    
     test_img = np.array(image)#cv2.imread(path, cv2.IMREAD_COLOR)       
     test_img = cv2.resize(test_img, (640, 640))
+    return test_img
     test_img = cv2.cvtColor(test_img, cv2.COLOR_RGB2BGR)
     test_img = np.expand_dims(test_img, axis=0)