import os
import sys

import torch
from torch.utils.serialization import load_lua
from torch.autograd import Variable

from vgg16 import Vgg16

import numpy as np
from PIL import Image
from scipy.misc import imread, imsave, imresize
from scipy.ndimage.filters import median_filter


# result: RGB CxHxW [0,255] torch.FloatTensor
def tensor_load_rgbimage(filename, size=None):
    img = Image.open(filename)
    if size is not None:
        img = img.resize((size, size), Image.ANTIALIAS)
    img = np.array(img).transpose(2, 0, 1)
    img = torch.from_numpy(img).float()
    return img


def gram_matrix(y):
    (b, ch, h, w) = y.size()
    features = y.view(b, ch, w * h)
    features_t = features.transpose(1, 2)
    gram = features.bmm(features_t) / (ch * h * w)
    return gram


def add_imagenet_mean(img):
    """Add ImageNet mean pixel-wise to a BGR image."""
    img[0, :, :] += 103.939
    img[1, :, :] += 116.779
    img[2, :, :] += 123.68


def deprocess_img_and_save(img, filename):
    """Undo pre-processing on an image, and save it."""
    img = img[0, :, :, :]
    add_imagenet_mean(img)
    img = img[::-1].transpose((1, 2, 0))
    img = np.clip(img, 0, 255).astype(np.uint8)
    img = median_filter(img, size=(3, 3, 1))
    try:
        imsave(filename, img)
    except OSError as e:
        print(e)
        sys.exit(1)


def subtract_imagenet_mean_batch(batch):
    """Subtract ImageNet mean pixel-wise from a BGR image."""
    tensortype = type(batch.data)
    mean = tensortype(batch.data.size())
    mean[:, 0, :, :] = 103.939
    mean[:, 1, :, :] = 116.779
    mean[:, 2, :, :] = 123.680
    batch -= Variable(mean)


def preprocess_batch(batch):
    batch = batch.transpose(0, 1)
    (r, g, b) = torch.chunk(batch, 3)
    batch = torch.cat((b, g, r))
    batch = batch.transpose(0, 1)
    return batch


def init_vgg16(model_folder):
    """load the vgg16 model feature"""
    if not os.path.exists(model_folder + '/vgg16.weight'):
        if not os.path.exists(model_folder + '/vgg16.t7'):
            os.system(
                'wget http://cs.stanford.edu/people/jcjohns/fast-neural-style/models/vgg16.t7 -O ' + model_folder + '/vgg16.t7')
        vgglua = load_lua(model_folder + '/vgg16.t7')
        vgg = Vgg16()
        for (src, dst) in zip(vgglua.parameters()[0], vgg.parameters()):
            dst[:].data = src[:]
        torch.save(vgg.state_dict(), model_folder + '/vgg16.weight')

@@ -182,104 +182,8 @@
 age%0A
-from scipy.misc import imread, imsave, imresize%0Afrom scipy.ndimage.filters import median_filter%0A
 %0A%0A# 
@@ -689,322 +689,99 @@
 def 
-add_imagenet_mean(img):%0A    %22%22%22Add ImageNet mean pixel-wise to a BGR image.%22%22%22%0A    img%5B0, :, :%5D += 103.939%0A    img%5B1, :, :%5D += 116.779%0A    img%5B2, :, :%5D += 123.68%0A%0A%0Adef deprocess_img_and_save(img, filename):%0A    %22%22%22Undo pre-processing on an image, and save it.%22%22%22%0A    img = img%5B0, :, :, :%5D%0A    add_imagenet_mean(img
+tensor_save_rgbimage(tensor, filename):%0A    img = tensor.clone().cpu().clamp(0, 255).numpy(
 )%0A  
@@ -795,14 +795,8 @@
  img
-%5B::-1%5D
 .tra
@@ -806,17 +806,16 @@
 ose(
-(
 1, 2, 0)
 )%0A  
@@ -814,48 +814,16 @@
 , 0)
-)%0A    img = np.clip(img, 0, 255)
 .astype(
 np.u
@@ -822,16 +822,15 @@
 ype(
-np.
+'
 uint8
+'
 )%0A  
@@ -841,142 +841,212 @@
 g = 
-median_filter(img, size=(3, 3, 1))%0A    try:%0A        imsave(filename, img)%0A    except OSError as e:%0A        print(e)%0A        sys.exit(1
+Image.fromarray(img)%0A    img.save(filename)%0A%0A%0Adef tensor_save_bgrimage(tensor, filename):%0A    (b, g, r) = torch.chunk(tensor, 3)%0A    tensor = torch.cat((r, g, b))%0A    tensor_save_rgbimage(tensor, filename
 )%0A%0A%0A
@@ -2036,19 +2036,16 @@
      dst
-%5B:%5D
 .data = 
@@ -2051,11 +2051,16 @@
  src
-%5B:%5D
+.clone()
 %0A   

Fix mistake in loading of vgg