projection only works with queue

Segmentation/segmentation.py

@@ -1,39 +1,33 @@
-# Import libraries
-
 import cv2
 from tensorflow import keras
 import numpy as np
-import matplotlib.pyplot as plt
 from PIL import Image
 import segmentation_models as sm
 sm.set_framework('tf.keras')
 
+# Load model at build time
 model_path = "Segmentation/model_checkpoint.h5"
 CLASSES = ['sofa']
 BACKBONE = 'resnet50'
 
-# define network parameters
-n_classes = 1 if len(CLASSES) == 1 else (len(CLASSES) + 1)  # case for binary and multiclass segmentation
+# define network parameters (only neede to load the weights)
+n_classes = 1 if len(CLASSES) == 1 else (len(CLASSES) + 1) 
 activation = 'sigmoid' if n_classes == 1 else 'softmax'
 preprocess_input = sm.get_preprocessing(BACKBONE)
 LR=0.0001
-
 #create model architecture
 model = sm.Unet(BACKBONE, classes=n_classes, activation=activation)
 # define optomizer
 optim = keras.optimizers.Adam(LR)
-# Segmentation models losses can be combined together by '+' and scaled by integer or float factor
 dice_loss = sm.losses.DiceLoss()
 focal_loss = sm.losses.BinaryFocalLoss() if n_classes == 1 else sm.losses.CategoricalFocalLoss()
 total_loss = dice_loss + (1 * focal_loss)
-# actulally total_loss can be imported directly from library, above example just show you how to manipulate with losses
-# total_loss = sm.losses.binary_focal_dice_loss # or sm.losses.categorical_focal_dice_loss 
 metrics = [sm.metrics.IOUScore(threshold=0.5), sm.metrics.FScore(threshold=0.5)]
 # compile keras model with defined optimozer, loss and metrics
 model.compile(optim, total_loss, metrics)
 model.load_weights(model_path) 
 
-def get_mask(image:Image) -> Image:
+def get_mask(image:Image.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.
@@ -57,7 +51,7 @@ def get_mask(image:Image) -> Image:
     mask = Image.fromarray(prediction[...,0].squeeze()*255).convert("L")
     return mask 
 
-def replace_sofa(image:Image, mask:Image, styled_sofa:Image) -> Image:
+def replace_sofa(image:Image.Image, mask:Image.Image, styled_sofa:Image.Image) -> Image.Image:
     """
     This function replaces the original sofa in the image by the new styled sofa according
     to the mask.
@@ -77,12 +71,3 @@ def replace_sofa(image:Image, mask:Image, styled_sofa:Image) -> Image:
     sofa_fg = cv2.bitwise_and(styled_sofa,styled_sofa,mask = mask)
     new_image = cv2.add(image_bg,sofa_fg)
     return Image.fromarray(new_image)
-
-# image = cv2.imread('input/sofa.jpg')
-# mask = cv2.imread('masks/sofa.jpg')
-# styled_sofa = cv2.imread('output/sofa_stylized_style.jpg')
-
-# #get_mask(image)
-
-# plt.imshow(replace_sofa(image,mask,styled_sofa))
-# plt.show()

StyleTransfer/styleTransfer.py

@@ -1,6 +1,7 @@
 from PIL import Image
 import numpy as np
 import torch
+torch.cuda.is_available()
 import torch.nn as nn
 from torchvision import transforms
 import StyleTransfer.transformer as transformer