Style Transfer (tutorial)

tutorial
computer-vision
neural-style-transfer

#1

Style Transfer

Bạn nghĩ sao về một bức ảnh chụp (máy ảnh) Hà Nội nhưng lại mang phong cách tranh của Bùi Xuân Phái :D. Với sự ra đời của thuật toán Style Transfer, chuyện đó là hoàn toàn có thể. Trên đây là sản phẩm của mình với thuật toán

Note: Để dễ đọc tutorial hơn, Mọi người có thể vào blog của mình, trong đó có vài bài khác khá thú vị đấy:

pre Knowledge

Bài trước, mình đã viết và code thuật toán Visualize Filter. Style Transfer có áp dụng 1 phần ý tưởng trong Visualize Filter, bạn nên đọc lại bài Visualize Filter mình đã viết tại:

OK, chúng ta bắt đầu thôi :slight_smile:

1. Thuật toán.

Note: đường update gradient mình vẽ thiếu mũi tên, chiều mũi tên từ combine_loss hướng tới generate_image

Trên đây là hình minh họa cho thuật toán. Chúng ta có 3 ảnh input gồm:

  • content_image: chứa nội dung mà generate_image sẽ chứa.
  • style_image: chửa style (phong cách) mà generate_image sẽ chứa
  • generate_image: Được khởi tạo random bất kì. sau quá trình update sẽ thành ảnh mong muốn.

Ý tưởng cơ bản

Ý tưởng thuật toán rất đơn giản: cả 3 ảnh cùng đưa vào 1 pretrain-CNN để extract ra các thông tin sau:

  • Với content_image thì extract ra content_feature (dữ liệu liên quan tới nội dung ảnh), ta gọi là content_A
  • Với style_image thì extract ra style_feature, gọi là style_B
  • Với generate_image thì extract ra cả 2 loại trên, gọi là content_Cstyle_C
  • Dựa vào 4 feature trên, ta tính “khoảng cách” hay “sự khác biệt” theo dạng: loss = (content_A - content_C) + (style_B - style_C)
  • Dựa vào loss để tính gradient của loss với biến là generate_image
  • Update generate_image với gradient
  • Quy trình sẽ lặp đi lặp lại nhằm tối ưu giá trị loss.

Content_feature

Câu hỏi đặt ra là làm thế nào để extract ra content_feature và style_feature, và làm thế nào để phân biệt giữa hai loại này? Với content_feature: output của các layer (các feature_map) trong CNN chính là content_feature, nó chứa thông tin về đường nét, bố cục, hình dạng, màu sắc của mọi thứ trong ảnh.

Style_feature

Thế còn Style_feature? cái này phức tạp hơn một chút. Cũng từ feature_map trên, nhưng ta phải biến đổi một chút ta mới nhận được style_feature:

  • Giả sử ta có feature_map A với shape: (1 * 5 * 5 * number_fmap)
  • Ta cần reshape A thành matrix 2 chiều: (25 * number_fmap). (bước này chính là flatten từng feature_map thành vector 1 chiều). Ta thu được matrix B với shape = (25 * feature_map).
  • Tính gram_matrix C = B^T * B
  • (B^T: ma trận chuyển vị của B, phép nhân là nhân ma trận).
  • Qua thực nghiệm, tác giả nhận ra rằng matrix C chính là thứ đặc trưng cho style của 1 bức ảnh (hoặc tranh).

Thuật toán về cơ bản là thế, bắt tay vào code bạn sẽ dễ hiểu hơn.

2. Code

Chuẩn bị:

  • Bạn nên vào github của mình và download code về tại: github.com/trungthanhnguyen0502/style-transfer.
  • Bạn có thể dùng bất kì pretrain CNN nào, tuy nhiên mình (và nhiều người) nhận thấy bài toán dạng này thì VGG là tốt nhất. Mình dùng 1 phiên bản VGG16 mà mình hay sử dụng. Để có thể sử dụng VGG cùng phiên bản với mình, hãy đảm bảo trong project của bạn có 2 file : download.pyvgg16.py
  • Mình sẽ chỉ hướng dẫn code chính, những việc lặt vặt như import thư viện, viết các function phụ, hãy đọc code của mình nhé.
  • 4 dòng đầu bạn có thể xóa đi nếu không sử dụng google colab.
  • Để đọc code dễ hiểu nhất, nên đọc từ hàm transfer để hiểu được kịch bản code. Trong transfer sẽ gọi đến các hàm khác. Mình viết tutorial theo đúng thứ tự code nhưng bạn không nên đọc theo thứ tự này, cứ đọc từ hàm transfer bên dưới nhé.
  • Có thể nhiều chỗ mình viết code không dễ đọc. Có gì khó hiểu, comment hoặc liên lạc qua fb: https://www.facebook.com/trungthanhnguyen0502. (Hãy gửi 1 lời trước khi gửi friend request nhé, nếu không mình tưởng là bán hàng online, mình lại không accept)

load và tạo model:

vgg16.maybe_download()
vgg = vgg16.VGG16()

hàm tính mean_square_error và tính gram_matrix

def mean_square_error(a,b):
    return tf.reduce_mean(tf.square(a-b))

def gram_matrix(a):    
    shape = a.get_shape()
    reshape_tensor = tf.reshape(a, shape=[-1, int(shape[3])])
    gram_tensor = tf.matmul(tf.transpose(reshape_tensor), reshape_tensor)
    return gram_tensor

content_loss

def content_loss(model, session, content_img, c_layer_ids):
    with model.graph.as_default():
        c_layers = model.get_layer_tensors(c_layer_ids)
        feed_dict = model.create_feed_dict(image=content_img)
        c_values = session.run(c_layers, feed_dict=feed_dict)
        c_loss = []
        for v, l in zip(c_values, c_layers):
            loss = mean_square_error(l, v)
            c_loss.append(loss)
        return tf.reduce_mean(c_loss)

style_loss

def style_loss(model, session, style_img, s_layer_ids):
    with model.graph.as_default():
        s_layers = model.get_layer_tensors(s_layer_ids)
        gram_matrix_tensor = [gram_matrix(l) for l in s_layers]
        feed_dict = model.create_feed_dict(image=style_img)
        gram_matrix_values = session.run(gram_matrix_tensor, feed_dict=feed_dict)
        s_loss = []
        for v, l in zip(gram_matrix_values, gram_matrix_tensor):
            loss = mean_square_error(v,l)
            s_loss.append(loss)
        return tf.reduce_mean(s_loss)

denoise_loss

denoise_loss là 1 kĩ thuật được bổ sung để generate_image giảm thiểu sự sai khác giữa 2 pixel cạnh nhau. (input_[:,1:,:,:] - input_[:,:-1,:,:])) và (input_[:,:,1:,:] - input_[:,:,:-1,:]) chính là tính độ sai khác giữa 2 pixel cạnh nhau theo chiều x và y

def denoise_loss(input_):
    d_loss = tf.reduce_mean(tf.abs(input_[:,1:,:,:] - input_[:,:-1,:,:])) + \
            tf.reduce_mean(tf.abs(input_[:,:,1:,:] - input_[:,:,:-1,:]))
    return d_loss

transfer: hàm cha, là kịch bản chính, gọi tới các hàm trên.

def transfer(model, content_img, style_img,
             c_layer_ids, s_layer_ids, w_content=1.5, 
             w_style=10.0, w_denoise=0.3, iters= 60, step_size=10):
    img_result = np.random.rand(*content_img.shape) + 128
    session = tf.Session(graph = model.graph)

    c_loss = content_loss(model, session, content_img, c_layer_ids)
    s_loss = style_loss( model, session, style_img, s_layer_ids)
    d_loss = denoise_loss(model.input)
    with model.graph.as_default():
        adj_content = tf.Variable(1e-10, name='adj_content')
        adj_style = tf.Variable(1e-10, name='adj_style')
        adj_denoise = tf.Variable(1e-10, name='adj_denoise')

        session.run([adj_content.initializer,
                     adj_style.initializer,
                     adj_denoise.initializer])
 
        update_adj_content = adj_content.assign(1/(c_loss + 1e-10))
        update_adj_style = adj_style.assign(1/(s_loss + 1e-10))
        update_adj_denoise = adj_denoise.assign(1/(d_loss + 1e-10))

        total_loss = w_content * adj_content * c_loss +  w_denoise * adj_denoise * d_loss + w_style * adj_style * s_loss
        grad = tf.gradients(total_loss, model.input)
        run_list = [grad, update_adj_content, update_adj_style, update_adj_denoise]
        
        for i in range(iters):
            print(i)
            feed_dict = model.create_feed_dict(image=img_result)
            grad_value, _, _, _ = session.run(run_list, feed_dict=feed_dict)
            grad_value = np.squeeze(grad_value)
            grad_value = grad_value.reshape(img_result.shape)
            
            scale_step = step_size/(np.std(grad_value) + 1e-8)
            img_result -= grad_value*scale_step
            img_result = np.clip(img_result, 0.0, 255.0)
            if( i % 10 == 0):
                plot_image(img_result)
    return img_result

Trong đó:

  • Tensorflow cung cấp cho chúng ta hàm tf.gradient(loss, tensor_variable) để tính đạo hàm loss theo tensor_variable.
  • Ta áp dụng thuật toán optimize đơn giản, learning rate: lrate = \frac{stepsize}{(gradient.max() + 1e-8)}
  • 1e-8: số cực đảm bảo mẫu số != 0
  • step_size là 1 giá trị cho trước
  • gradient.max() --> giá trị update của 1 pixel: update_value_pixel[x,y,z] <= step_size* gradient_value[x,y,z] (gradient_value tương ứng pixel) . Ngoài gradient.max(), ta có thể dùng gradient.std()

Tiến hành thử nghiệm:

content_dir = 'data/content.jpg'
content_img = load_image(content_dir)
content_img = cv2.resize(content_img, (224,224))

style_dir = 'data/style.jpeg'
style_img = load_image(style_dir)
style_img = cv2.resize(style_img, (224,224))

c_layer_ids = [4]
s_layer_ids = range(13)
result_img = transfer(model, content_img, style_img, 
                      c_layer_ids, s_layer_ids, 
                      w_content=2.0, w_style=9.0, w_denoise=0.4, 
                      iters= 60, step_size=10)

3 trọng số: w_content, w_style, w_denoise sau khi mình run code nhiều lần thấy w_content giao động quanh giá trị 2, w_style ~ 10, w_denoise ~ 0.5 sẽ cho kết quả tốt nhất. Kết qủa: với các tham số mình đã đặt mặc định, chỉ sau 60 vòng lặp, ta đã có thể thu được kết quả mong muốn:

Bạn có thể tự custom code, như điều chỉnh tham số, bỏ denoise_loss, áp dụng với ảnh khác … Chúc bạn thành công nhé.

Liên lạc với tôi:


#2

Mình đã thử.