Fast Neural Style Transfer in Owl

#readme.md

Fast Neural Style Transfer

Neural Style Transfer is the process of using Deep Neural Networks to migrate the semantic content of one image to different styles. Fast Neural Style Transfer (FST) can finish this process in the order of seconds.

Usage

This gist implements Fast Style Transfer in Owl, and provides a simple interfaces to use. Here is an example:

#zoo "f937ce439c8adcaea23d42753f487299"

FST.list_styles ();; (* show all supported styles *)
FST.run ~style:1 "path/to/content_img.png" "path/to/output_img.jpg" 

The run function mainly takes one content image and output to a new image file, the name of which is designated by the user. The image chould be of any popular formats: jpeg, png, etc. This gist contains exemplar content images for you to use.

Current we support six art styles:

1. "Udnie" by Francis Picabia
2. "The Great Wave off Kanagawa" by Hokusai
3. "Rain Princess" by Leonid Afremov
4. "La Muse" by Picasso
5. "The Scream" by Edvard Munch
6. "The shipwreck of the Minotaur" by J. M. W. Turner

Prerequisite

This application relies on the tool ImageMagick to manipulate image format conversion and resizing. Please make sure it is installed. E.g. on Ubuntu:

sudo apt-get install imagemagick

Limit

FST is limited in styles. If you want to try some styles that are not listed here, please refer to our Neural Style Transfer gist.

example.png

FST.ml

open Owl
open Neural.S
open Neural.S.Graph
open Neural.S.Algodiff
module N = Dense.Ndarray.S

#zoo "86a1748bbc898f2e42538839edba00e1" (* ImageUtils  *)

(** Network Structure *)

let conv2d_layer ?(relu=true) kernel stride nn  =
  let result = 
    conv2d ~padding:SAME kernel stride nn
    |> normalisation ~decay:0. ~training:true ~axis:3
  in
  match relu with
  | true -> (result |> activation Activation.Relu)
  | _    -> result

let conv2d_trans_layer kernel stride nn = 
  transpose_conv2d ~padding:SAME kernel stride nn
  |> normalisation ~decay:0. ~training:true ~axis:3
  |> activation Activation.Relu

let residual_block wh nn = 
  let tmp = conv2d_layer [|wh; wh; 128; 128|] [|1;1|] nn
    |> conv2d_layer ~relu:false [|wh; wh; 128; 128|] [|1;1|]
  in 
  add [|nn; tmp|]

(* perfectly balanced -- like everything should be. *)
let make_network h w = 
  input [|h;w;3|]
  |> conv2d_layer [|9;9;3;32|] [|1;1|]
  |> conv2d_layer [|3;3;32;64|] [|2;2|]
  |> conv2d_layer [|3;3;64;128|] [|2;2|]
  |> residual_block 3
  |> residual_block 3
  |> residual_block 3
  |> residual_block 3
  |> residual_block 3
  |> conv2d_trans_layer [|3;3;128;64|] [|2;2|]
  |> conv2d_trans_layer [|3;3;64;32|] [|2;2|]
  |> conv2d_layer ~relu:false [|9;9;32;3|] [|1;1|]
  |> lambda (fun x -> Maths.((tanh x) * (F 150.) + (F 127.5)))
  |> get_network

(* Image helper functions *)
let _convert img_name = 
  let base = Filename.basename img_name in 
  let prefix = Filename.remove_extension base in
  let temp_img = Filename.temp_file prefix ".ppm"in
  temp_img

let convert_img_to_ppm w h img_name = 
  let temp_img = _convert img_name in
  let _ = Sys.command ("convert -resize " ^ (string_of_int w) ^ 
    "x" ^ (string_of_int h) ^"\\! " ^ 
    img_name ^ " " ^ temp_img) in
  temp_img

let convert_arr_to_img d3array output_name = 
  let temp_img = _convert output_name in
  let output = d3array in
  ImageUtils.save_ppm_from_arr output temp_img;
  let _ = Sys.command ("convert " ^ temp_img ^ " " ^ output_name) in
  ()

let get_img_shape img_name = 
  let temp_img = _convert img_name in
  let _ = Sys.command ("convert " ^ img_name ^ " " ^ temp_img) in
  let _, w, h, _ = ImageUtils._read_ppm temp_img in
  w, h

(* Styles *)

let styles = [|"udnie"; "wave"; "rain_princess"; "la_muse"; "scream"; "wreck"|] 

let make_style_htb () = 
  let h = Hashtbl.create 10 in
  for i = 0 to (Array.length styles - 1) do
    (* weight file: e.g. fst_udnie.weight *)
    Hashtbl.add h i ("fst_" ^ styles.(i) ^ ".weight")
  done;
  h

let style_htb = make_style_htb ()  

let list_styles () = 
  let s = ref "" in
  for i = 0 to (Array.length styles - 1) do
    s := !s ^ Printf.sprintf "Style [%d] :\t %s\n" i styles.(i)
  done; 
  let info = Printf.sprintf "Here are the usable styles:\n%s" !s in
  print_endline info

(* FST service function *)
let run ?(style=0) content_img output_img =
  
  let w, h = get_img_shape content_img in
  let content_img = convert_img_to_ppm w h content_img in
  let content_img = ImageUtils.(load_ppm content_img |> extend_dim) in

  let nn = make_network h w in
  Graph.init nn;

  let style_file = 
    try Hashtbl.find style_htb style 
    with Not_found -> failwith "style does not exist; try to run `list_styles ()`"
  in
  Graph.load_weights nn (Owl_zoo_path.extend_zoo_path style_file);
  let result = Graph.model nn content_img in

  convert_arr_to_img result output_img

fst_la_muse.weight

fst_rain_princess.weight

fst_scream.weight

fst_udnie.weight

fst_wave.weight

fst_wreck.weight

Here shows the FST results for the example image with different styles.

"We're in the endgame now"