make transformer global

app.py

@@ -90,10 +90,12 @@ def style_sofa(input_img: np.ndarray, style_img: np.ndarray):
     input_img,style_img = Image.fromarray(input_img),Image.fromarray(style_img)
     resized_img,box = resize_sofa(input_img)
     resized_style = resize_style(style_img)
+    resized_style.save('resized_style.jpg')
     # generate mask for image
     mask = get_mask(resized_img)
+    mask.save('mask.jpg')
     styled_sofa = create_styledSofa(resized_img,resized_style)
-    print(type(styled_sofa))
+    styled_sofa.save('styled_sofa.jpg')
     # postprocess the final image
     new_sofa = replace_sofa(resized_img,mask,styled_sofa)
     new_sofa = new_sofa.crop(box)

styleTransfer.py

@@ -11,10 +11,15 @@ from collections import OrderedDict
 import tensorflow_hub as hub
 import tensorflow as tf
 
-style_transfer_model = hub.load("https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2")
-
+from torchvision.utils import save_image
 
 ############################################# TRANSFORMER ############################################
+
+vgg_path        = 'vgg_normalised.pth'
+decoder_path    = 'decoder_iter_160000.pth'
+Trans_path      = 'transformer_iter_160000.pth'
+embedding_path  = 'embedding_iter_160000.pth'
+
 def test_transform(size, crop):
     transform_list = []
    
@@ -40,82 +45,75 @@ def content_transform():
     transform = transforms.Compose(transform_list)
     return transform
 
-def StyleTransformer(content_img: Image, style_img: Image,
-                vgg_path:str = 'vgg_normalised.pth', decoder_path:str = 'decoder_iter_160000.pth',
-                Trans_path:str = 'transformer_iter_160000.pth', embedding_path:str = 'embedding_iter_160000.pth'):
-    # Advanced options
-    content_size=640
-    style_size=640
-    crop='store_true'
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    vgg = StyTR.vgg
-    vgg.load_state_dict(torch.load(vgg_path))
-    vgg = nn.Sequential(*list(vgg.children())[:44])
-
-    decoder = StyTR.decoder
-    Trans = transformer.Transformer()
-    embedding = StyTR.PatchEmbed()
-
-    decoder.eval()
-    Trans.eval()
-    vgg.eval()
-
-    new_state_dict = OrderedDict()
-    state_dict = torch.load(decoder_path)
-    for k, v in state_dict.items():
-        #namekey = k[7:] # remove `module.`
-        namekey = k
-        new_state_dict[namekey] = v
-    decoder.load_state_dict(new_state_dict)
-
-    new_state_dict = OrderedDict()
-    state_dict = torch.load(Trans_path)
-    for k, v in state_dict.items():
-        #namekey = k[7:] # remove `module.`
-        namekey = k
-        new_state_dict[namekey] = v
-    Trans.load_state_dict(new_state_dict)
-
-    new_state_dict = OrderedDict()
-    state_dict = torch.load(embedding_path)
-    for k, v in state_dict.items():
-        #namekey = k[7:] # remove `module.`
-        namekey = k
-        new_state_dict[namekey] = v
-    embedding.load_state_dict(new_state_dict)
-
-    network = StyTR.StyTrans(vgg,decoder,embedding,Trans)
-    network.eval()
-    network.to(device)
-
-    content_tf = test_transform(content_size, crop)
-    style_tf = test_transform(style_size, crop)
-    
+# Advanced options
+content_size=640
+style_size=640
+crop='store_true'
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+vgg = StyTR.vgg
+vgg.load_state_dict(torch.load(vgg_path))
+vgg = nn.Sequential(*list(vgg.children())[:44])
+
+decoder = StyTR.decoder
+Trans = transformer.Transformer()
+embedding = StyTR.PatchEmbed()
+
+decoder.eval()
+Trans.eval()
+vgg.eval()
+
+new_state_dict = OrderedDict()
+state_dict = torch.load(decoder_path)
+for k, v in state_dict.items():
+    #namekey = k[7:] # remove `module.`
+    namekey = k
+    new_state_dict[namekey] = v
+decoder.load_state_dict(new_state_dict)
+
+new_state_dict = OrderedDict()
+state_dict = torch.load(Trans_path)
+for k, v in state_dict.items():
+    #namekey = k[7:] # remove `module.`
+    namekey = k
+    new_state_dict[namekey] = v
+Trans.load_state_dict(new_state_dict)
+
+new_state_dict = OrderedDict()
+state_dict = torch.load(embedding_path)
+for k, v in state_dict.items():
+    #namekey = k[7:] # remove `module.`
+    namekey = k
+    new_state_dict[namekey] = v
+embedding.load_state_dict(new_state_dict)
+
+network = StyTR.StyTrans(vgg,decoder,embedding,Trans)
+network.eval()
+network.to(device)
+
+content_tf = test_transform(content_size, crop)
+style_tf = test_transform(style_size, crop)
+
+def StyleTransformer(content_img: Image, style_img: Image):
     content_tf1 = content_transform()       
     content = content_tf(content_img.convert("RGB"))
-
     h,w,c=np.shape(content)    
     style_tf1 = style_transform(h,w)
     style = style_tf(style_img.convert("RGB"))
-
-
     style = style.to(device).unsqueeze(0)
     content = content.to(device).unsqueeze(0)
 
     with torch.no_grad():
         output= network(content,style)  
-    print(type(output))  
-    output = output[0].cpu()
-    print(type(output))
-    print(output.squeeze().shape)
-    torch2PIL = transforms.ToPILImage()
-    output = torch2PIL(output.squeeze())
-    return output
+    output = output[0].cpu().squeeze()
+    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_transfer_model = hub.load("https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2")
+
 def perform_style_transfer(content_image, style_image):
     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.
@@ -130,10 +128,7 @@ def create_styledSofa(sofa:Image, style:Image):
     styled_sofa = StyleTransformer(sofa,style)
     return styled_sofa
 
-# image = Image.open('sofa_office.jpg') 
-# image.show() 
-# image = np.array(image)
-# image,box = resize_sofa(image)  
-# image = image.crop(box)
-# print(box)
-# image.show()
+image = Image.open('sofa.jpg') 
+style = Image.open('style.jpg')
+output = create_styledSofa(image,style)
+output.show()