Style-based GAN
File.cd!(__DIR__)
# for windows JP
System.shell("chcp 65001")
System.put_env("NNCOMPILED", "YES")
Mix.install([
{:tfl_interp, "~> 0.1.13"},
{:cimg, "~> 0.1.19"},
{:kino, "~> 0.7.0"},
{:nx, "~> 0.5.3"}
])
0.Original work
“Training Generative Adversarial Networks with Limited Data”
GitHub: StyleGAN2 with adaptive discriminator augmentation
GitHub: sugyan/gan-playground
The tflite model afhqdog.tflite
is converted from pretraind model in pickle format.
Thanks a lot!!!
Implementation with TflInterp in Elixir
1.Defining the inference module: StyleGAN2ADA
-
Model
afhqdog.tflite: get from “https://github.com/shoz-f/tfl_interp/releases/download/0.0.1/afhqdog.tflite” if not existed.
-
Pre-processing
Make a latent - 512 size normal random number vector.
-
Post-processing
Convert float32 {-1.0, 1.0} to pixel and transpose NCHW to NHWC.
defmodule StyleGAN2ADA.Mapping do
alias TflInterp, as: NNInterp
use NNInterp,
model: "./model/afhqdog.mapping.tflite",
url: "https://github.com/shoz-f/tfl_interp/releases/download/0.0.1/afhqdog.mapping.tflite",
inputs: [f32: {1, 512}],
outputs: [f32: {1, 16.512}]
def dlatants() do
for _ <- 1..512, into: "" do
<<:rand.normal()::little-float-32>>
end
|> Nx.from_binary(:f32)
|> Nx.reshape({1, 512})
|> dlatants()
end
def dlatants(latants) do
# preprocess
input0 = Nx.to_binary(latants)
# prediction
output0 =
session()
|> NNInterp.set_input_tensor(0, input0)
|> NNInterp.invoke()
|> NNInterp.get_output_tensor(0)
# postprocess
Nx.from_binary(output0, :f32) |> Nx.reshape({16, 512})
end
end
defmodule StyleGAN2ADA.Synthesis do
@width 512
@height 512
alias TflInterp, as: NNInterp
use NNInterp,
model: "./model/afhqdog.synthesis.tflite",
url: "https://github.com/shoz-f/tfl_interp/releases/download/0.0.1/afhqdog.synthesis.tflite",
inputs: [f32: {1, 16, 512}],
outputs: [f32: {1, 3, @height, @width}]
def image(dlatants) do
# preprocess
input0 = Nx.to_binary(dlatants)
# prediction
output0 =
session()
|> NNInterp.set_input_tensor(0, input0)
|> NNInterp.invoke()
|> NNInterp.get_output_tensor(0)
# postprocess
CImg.from_binary(output0, @width, @height, 1, 3, [{:range, {-1.0, 1.0}}, :nchw])
end
end
Launch StyleGAN2ADA.Mapping
and StyleGAN2ADA.Synthesis
.
# TflInterp.stop(StyleGAN2ADA.Mapping)
StyleGAN2ADA.Mapping.start_link([])
# TflInterp.stop(StyleGAN2ADA.Synthesis)
StyleGAN2ADA.Synthesis.start_link([])
Display the properties of the StyleGAN2ADA
model.
TflInterp.info(StyleGAN2ADA.Mapping)
TflInterp.info(StyleGAN2ADA.Synthesis)
2.Let’s try it
Morphing between two randomly generated images.
a = StyleGAN2ADA.Mapping.dlatants()
b = StyleGAN2ADA.Mapping.dlatants()
delta = Nx.subtract(b, a) |> Nx.divide(5)
Enum.map(0..5, fn i ->
Nx.add(a, Nx.multiply(delta, i))
|> StyleGAN2ADA.Synthesis.image()
|> CImg.display_kino(:jpeg)
end)
|> Kino.Layout.grid(columns: 6)
□