fix errors

styleTransfer.py

@@ -11,8 +11,6 @@ from collections import OrderedDict
 import tensorflow_hub as hub
 import tensorflow as tf
 
-from torchvision.utils import save_image
-
 ############################################# TRANSFORMER ############################################
 
 vgg_path        = 'vgg_normalised.pth'
@@ -20,16 +18,6 @@ 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 = []
-   
-    if size != 0: 
-        transform_list.append(transforms.Resize(size))
-    if crop:
-        transform_list.append(transforms.CenterCrop(size))
-    transform_list.append(transforms.ToTensor())
-    transform = transforms.Compose(transform_list)
-    return transform
 def style_transform(h,w):
     k = (h,w)
     size = int(np.max(k))
@@ -48,7 +36,6 @@ def content_transform():
 # Advanced options
 content_size=640
 style_size=640
-crop='store_true'
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -66,42 +53,28 @@ 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)
+decoder.load_state_dict(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)
+Trans.load_state_dict(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)
+embedding.load_state_dict(state_dict)
 
 network = StyTR.StyTrans(vgg,decoder,embedding,Trans)
 network.eval()
 
+content_tf = content_transform() 
+style_tf   = style_transform(style_size,style_size)
+
 def StyleTransformer(content_img: Image, style_img: Image):
    
     network.to(device)
-    # content_tf = test_transform(content_size, crop)
-    # style_tf = test_transform(style_size, crop)
-
-    content_tf1 = content_transform()       
-    content = content_tf1(content_img.convert("RGB"))
-    h,w,c=np.shape(content)    
-    style_tf1 = style_transform(h,w)
-    style = style_tf1(style_img.convert("RGB"))
+      
+    content = content_tf(content_img.convert("RGB"))   
+    style = style_tf(style_img.convert("RGB"))
     style = style.to(device).unsqueeze(0)
     content = content.to(device).unsqueeze(0)
 
@@ -128,4 +101,4 @@ def StyleGAN(content_image, style_image):
 def create_styledSofa(sofa:Image, style:Image):
     #styled_sofa = StyleGAN(sofa,style)
     styled_sofa = StyleTransformer(sofa,style)
-    return styled_sofa
+    return styled_sofa