styleTransfer.py

import torch
from PIL import Image
import numpy as np
import os
import cv2

def resize_sofa(img):
    img = Image.fromarray(img)
    width, height = img.size
    idx = np.argmin([width,height])

    if idx==0:
        img1 = Image.new(img.mode, (height, height), (255, 255, 255))
        img1.paste(img, ((height-width)//2, 0))
    else:
        img1 = Image.new(img.mode, (width, width), (255, 255, 255))
        img1.paste(img, (0, (width-height)//2))

    newsize = (640, 640) # parameters from test script
    im1 = img1.resize(newsize)
    return im1

def resize_style(img):
    #img = Image.open(path)#"../style5.jpg")
    img = Image.fromarray(img)
    width, height = img.size
    idx = np.argmin([width,height])
    #print(width,height)

    if idx==0:
        top= (height-width)//2
        bottom= height-(height-width)//2
        left = 0
        right= width
    else:
        left = (width-height)//2
        right = width - (width-height)//2
        top = 0
        bottom = height

    newsize = (640, 640) # parameters from test script
    im1 = img.crop((left, top, right, bottom))

    copies = 8
    resize = (newsize[0]//copies,newsize[1]//copies)
    dst = Image.new('RGB', (resize[0]*copies,resize[1]*copies))
    im2 = im1.resize((resize))
    for row in range(copies):
        im2 = im2.transpose(Image.FLIP_LEFT_RIGHT)
        for column in range(copies):
            im2 = im2.transpose(Image.FLIP_TOP_BOTTOM)
            dst.paste(im2, (resize[0]*row, resize[1]*column))
    dst = dst.resize((newsize))
    return dst

def create_styledSofa(sofa,style):
    path_sofa,path_style = 'sofa.jpg','style.jpg'
    sofa.save(path_sofa)
    style.save(path_style)
    os.system("python3 test.py  --content "+path_sofa+" \
                    --style "+path_style+" \
                    --output . \
                    --vgg vgg_normalised.pth \
                    --decoder_path decoder_iter_160000.pth \
                    --Trans_path transformer_iter_160000.pth \
                    --embedding_path embedding_iter_160000.pth")
    styled_sofa = cv2.imread('sofa_stylized_style.jpg')

    return styled_sofa

# image = Image.open('input/sofa.jpg')  
# image = np.array(image)[:,:600]
# image = resize_sofa(image)