forgot to update a file

StyleTransfer/styleTransfer.py

@@ -13,7 +13,7 @@ import paddlehub as phub
 
 ############################################# TRANSFORMER ############################################
 
-def style_transform(h,w):
+def style_transform(h:int,w:int) -> transforms.Compose:
     k = (h,w)
     transform_list = []    
     transform_list.append(transforms.CenterCrop((h,w)))
@@ -21,13 +21,13 @@ def style_transform(h,w):
     transform = transforms.Compose(transform_list)
     return transform
 
-def content_transform():
+def content_transform() -> transforms.Compose:
     transform_list = []   
     transform_list.append(transforms.ToTensor())
     transform = transforms.Compose(transform_list)
     return transform
 
-def StyleTransformer(content_img: Image, style_img: Image):
+def StyleTransformer(content_img: Image.Image, style_img: Image.Image) -> Image.Image:
     vgg_path        = 'StyleTransfer/models/vgg_normalised.pth'
     decoder_path    = 'StyleTransfer/models/decoder_iter_160000.pth'
     Trans_path      = 'StyleTransfer/models/transformer_iter_160000.pth'
@@ -43,7 +43,6 @@ def StyleTransformer(content_img: Image, style_img: Image):
     decoder = StyTR.decoder
     Trans = transformer.Transformer()
     embedding = StyTR.PatchEmbed()
-
     decoder.eval()
     Trans.eval()
     vgg.eval()
@@ -62,7 +61,6 @@ def StyleTransformer(content_img: Image, style_img: Image):
 
     network = StyTR.StyTrans(vgg,decoder,embedding,Trans)
     network.eval()
-
     content_tf = content_transform() 
     style_tf   = style_transform(style_size,style_size)
 
@@ -78,23 +76,20 @@ def StyleTransformer(content_img: Image, style_img: Image):
     output = output.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
     return Image.fromarray(output)
    
-############################################## STYLE-GAN #############################################
-
+############################################## STYLE-FAST #############################################
 
-def StyleGAN(content_image, style_image):
+def StyleFAST(content_image:Image.Image, style_image:Image.Image) -> Image.Image:
     style_transfer_model = tfhub.load("https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2")
 
-    content_image = tf.convert_to_tensor(content_image, np.float32)[tf.newaxis, ...] / 255.
-    style_image = tf.convert_to_tensor(style_image, np.float32)[tf.newaxis, ...] / 255.
+    content_image = tf.convert_to_tensor(np.array(content_image), np.float32)[tf.newaxis, ...] / 255.
+    style_image = tf.convert_to_tensor(np.array(style_image), np.float32)[tf.newaxis, ...] / 255.
     output = style_transfer_model(content_image, style_image)
     stylized_image = output[0]
     return Image.fromarray(np.uint8(stylized_image[0] * 255))
 
 ########################################### STYLE PROJECTION ##########################################
 
-
-
-def styleProjection(content_image,style_image):
+def StyleProjection(content_image:Image.Image,style_image:Image.Image) -> Image.Image:
     stylepro_artistic = phub.Module(name="stylepro_artistic")
     result = stylepro_artistic.style_transfer(
     images=[{
@@ -104,9 +99,15 @@ def styleProjection(content_image,style_image):
 
     return Image.fromarray(np.uint8(result[0]['data'])[:,:,::-1]).convert('RGB')
 
-
-def create_styledSofa(content_image,style_image):
-    output = StyleTransformer(content_image,style_image)
+def create_styledSofa(content_image:Image.Image,style_image:Image.Image,choice:str) -> Image.Image:
+    if choice =="Style Transformer":
+        output = StyleTransformer(content_image,style_image)
+    elif choice =="Style FAST":
+        output = StyleFAST(content_image,style_image)
+    elif choice =="Style Projection":
+        output = StyleProjection(content_image,style_image)
+    else:
+        output = content_image
     return output