Fast style transfer example

.gitignore

@@ -8,3 +8,4 @@
 .swiftpm
 cifar-10-batches-py/
 cifar-10-batches-bin/
+.idea

FastStyleTransfer/Demo/Helpers.swift

@@ -0,0 +1,92 @@
+import Foundation
+import TensorFlow
+import FastStyleTransfer
+
+// Make model importable
+extension TransformerNet: ImportableLayer {}
+
+func parseArgs<T>(into obj: inout T, with params: [String: WritableKeyPath<T, String?>]) {
+    for arg in CommandLine.arguments.dropFirst() {
+        if !arg.starts(with: "--") { continue }
+        let parts = arg.split(separator: "=", maxSplits: 2)
+        let name = String(parts[0][parts[0].index(parts[0].startIndex, offsetBy: 2)...])
+        if let path = params[name], parts.count == 2 {
+            obj[keyPath: path] = String(parts[1])
+        }
+    }
+}
+
+enum FileError: Error {
+    case file_not_found
+}
+
+func importWeights(_ model: inout TransformerNet, from path: String) throws {
+    guard FileManager.default.fileExists(atPath: path) else {
+        throw FileError.file_not_found
+    }
+    // Map of model params to loaded params
+    // Names don't match exactly, and axes in filters need to be reversed
+    let map = [
+        "conv1.conv2d.filter": ("conv1.conv2d.weight", [3, 2, 1, 0]),
+        "conv2.conv2d.filter": ("conv2.conv2d.weight", [3, 2, 1, 0]),
+        "conv3.conv2d.filter": ("conv3.conv2d.weight", [3, 2, 1, 0]),
+        "deconv1.conv2d.filter": ("deconv1.conv2d.weight", [3, 2, 1, 0]),
+        "deconv2.conv2d.filter": ("deconv2.conv2d.weight", [3, 2, 1, 0]),
+        "deconv3.conv2d.filter": ("deconv3.conv2d.weight", [3, 2, 1, 0]),
+        "res1.conv1.conv2d.filter": ("res1.conv1.conv2d.weight", [3, 2, 1, 0]),
+        "res1.conv2.conv2d.filter": ("res1.conv2.conv2d.weight", [3, 2, 1, 0]),
+        "res1.in1.scale": ("res1.in1.weight", nil),
+        "res1.in1.offset": ("res1.in1.bias", nil),
+        "res1.in2.scale": ("res1.in2.weight", nil),
+        "res1.in2.offset": ("res1.in2.bias", nil),
+        "res2.conv1.conv2d.filter": ("res2.conv1.conv2d.weight", [3, 2, 1, 0]),
+        "res2.conv2.conv2d.filter": ("res2.conv2.conv2d.weight", [3, 2, 1, 0]),
+        "res2.in1.scale": ("res2.in1.weight", nil),
+        "res2.in1.offset": ("res2.in1.bias", nil),
+        "res2.in2.scale": ("res2.in2.weight", nil),
+        "res2.in2.offset": ("res2.in2.bias", nil),
+        "res3.conv1.conv2d.filter": ("res3.conv1.conv2d.weight", [3, 2, 1, 0]),
+        "res3.conv2.conv2d.filter": ("res3.conv2.conv2d.weight", [3, 2, 1, 0]),
+        "res3.in1.scale": ("res3.in1.weight", nil),
+        "res3.in1.offset": ("res3.in1.bias", nil),
+        "res3.in2.scale": ("res3.in2.weight", nil),
+        "res3.in2.offset": ("res3.in2.bias", nil),
+        "res4.conv1.conv2d.filter": ("res4.conv1.conv2d.weight", [3, 2, 1, 0]),
+        "res4.conv2.conv2d.filter": ("res4.conv2.conv2d.weight", [3, 2, 1, 0]),
+        "res4.in1.scale": ("res4.in1.weight", nil),
+        "res4.in1.offset": ("res4.in1.bias", nil),
+        "res4.in2.scale": ("res4.in2.weight", nil),
+        "res4.in2.offset": ("res4.in2.bias", nil),
+        "res5.conv1.conv2d.filter": ("res5.conv1.conv2d.weight", [3, 2, 1, 0]),
+        "res5.conv2.conv2d.filter": ("res5.conv2.conv2d.weight", [3, 2, 1, 0]),
+        "res5.in1.scale": ("res5.in1.weight", nil),
+        "res5.in1.offset": ("res5.in1.bias", nil),
+        "res5.in2.scale": ("res5.in2.weight", nil),
+        "res5.in2.offset": ("res5.in2.bias", nil),
+        "in1.scale": ("in1.weight", nil),
+        "in1.offset": ("in1.bias", nil),
+        "in2.scale": ("in2.weight", nil),
+        "in2.offset": ("in2.bias", nil),
+        "in3.scale": ("in3.weight", nil),
+        "in3.offset": ("in3.bias", nil),
+        "in4.scale": ("in4.weight", nil),
+        "in4.offset": ("in4.bias", nil),
+        "in5.scale": ("in5.weight", nil),
+        "in5.offset": ("in5.bias", nil),
+    ]
+    model.unsafeImport(fromNumpyArchive: path, map: map)
+}
+
+func loadJpegAsTensor(from file: String) throws -> Tensor<Float> {
+    guard FileManager.default.fileExists(atPath: file) else {
+        throw FileError.file_not_found
+    }
+    let imgData = Raw.readFile(filename: StringTensor(file))
+    return Tensor<Float>(Raw.decodeJpeg(contents: imgData, channels: 3, dctMethod: "")) / 255
+}
+
+func saveTensorAsJpeg(_ tensor: Tensor<Float>, to file: String) {
+    let clipped = Raw.clipByValue(t: tensor, clipValueMin: Tensor(0), clipValueMax: Tensor(255))
+    let jpg = Raw.encodeJpeg(image: Tensor<UInt8>(clipped), format: .rgb, xmpMetadata: "")
+    Raw.writeFile(filename: StringTensor(file), contents: jpg)
+}

FastStyleTransfer/Demo/main.swift

@@ -0,0 +1,50 @@
+import Foundation
+import TensorFlow
+import FastStyleTransfer
+
+func printUsage() {
+    let exec = CommandLine.arguments[0].lastPathComponent
+    print("Usage:")
+    print("\(exec) --weights=<path> --image=<path> --output=<path>")
+    print("    --weights: Path to weights in numpy format")
+    print("    --image: Path to image in JPEG format")
+    print("    --output: Path to output image")
+}
+
+struct Config {
+    var weights: String? = "Demo/weights/candy.npz"
+    var image: String? = nil
+    var output: String? = "out.jpg"
+}
+var config = Config()
+parseArgs(into: &config, with: [
+    "weights": \Config.weights, 
+    "image": \Config.image, 
+    "output": \Config.output])
+
+guard let image = config.image, let output = config.output else {
+    print("Error: No input image!")
+    printUsage()
+    exit(1)
+}
+
+// load image
+guard let imageTensor = try? loadJpegAsTensor(from: image) else {
+    print("Error: Failed to load image \(image). Check file exists and has JPEG format")
+    printUsage()
+    exit(1)
+}
+
+// init model
+var style = TransformerNet()
+do {
+    try importWeights(&style, from: config.weights!)
+} catch {
+    print("Error: Failed to load weights \(config.weights!). Check file exists and has NPZ format")
+    printUsage()
+    exit(1)
+}
+
+let out = style(imageTensor.expandingShape(at: 0))
+saveTensorAsJpeg(out.squeezingShape(at: 0), to: output)
+print("Written output to \(output)")

FastStyleTransfer/Demo/weights/torch-convert.py

@@ -0,0 +1,15 @@
+import sys
+import torch
+import numpy as np
+
+# Usage:
+# python torch-convert.py model.pth model
+# (produces model.npz)
+
+if __name__ == "__main__":
+	in_file, out_file= sys.argv[1], sys.argv[2]
+	state_dict = torch.load(in_file)
+	npz = {}
+	for label, tensor in state_dict.items():
+		npz[label] = tensor.numpy()
+	np.savez(out_file, **npz)

FastStyleTransfer/Layers/Helpers.swift

@@ -0,0 +1,34 @@
+import TensorFlow
+
+/// Layer for padding with reflection over mini-batch of images
+/// Expected input layout is BxHxWxC
+public struct ReflectionPad2d<Scalar: TensorFlowFloatingPoint>: Layer {
+    @noDerivative public let padding: ((Int, Int), (Int, Int))
+
+    public init(padding: ((Int, Int), (Int, Int))) {
+        self.padding = padding
+    }
+
+    public init(padding: Int) {
+        self.padding = ((padding, padding), (padding, padding))
+    }
+
+    @differentiable
+    public func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
+        return input.paddedWithReflection(forSizes: [
+            (0, 0),
+            padding.0,
+            padding.1,
+            (0, 0)
+        ])
+    }
+}
+
+
+/// Layer applying relu activation function
+public struct ReLU<Scalar: TensorFlowFloatingPoint>: Layer {
+    @differentiable
+    public func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
+        return relu(input)
+    }
+}

FastStyleTransfer/Layers/Normalization.swift

@@ -0,0 +1,24 @@
+import TensorFlow
+
+/// Layer that applies instance normalization over a mini-batch of images
+/// Expected input layout is BxHxWxC
+/// Reference: [Instance Normalization](https://arxiv.org/abs/1607.08022)
+public struct InstanceNorm2d<Scalar: TensorFlowFloatingPoint>: Layer {
+    public var scale: Tensor<Scalar>
+    public var offset: Tensor<Scalar>
+    @noDerivative public var epsilon: Tensor<Scalar>
+
+    public init(featureCount: Int, epsilon: Tensor<Scalar> = Tensor(1e-5)) {
+        self.epsilon = epsilon
+        scale = Tensor<Scalar>(ones: [featureCount])
+        offset = Tensor<Scalar>(zeros: [featureCount])
+    }
+
+    @differentiable
+    public func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
+        let mean = input.mean(alongAxes: [1, 2])
+        let variance = input.variance(alongAxes: [1, 2])
+        let norm = (input - mean) * rsqrt(variance + epsilon)
+        return norm * scale + offset
+    }
+}

FastStyleTransfer/Models/TransformerNet.swift

@@ -0,0 +1,118 @@
+import TensorFlow
+import Foundation
+
+/// Model that applies style
+public struct TransformerNet: Layer {
+    public typealias Input = Tensor<Float>
+    public typealias Output = Tensor<Float>
+
+    // conv layers
+    public var conv1 = ConvLayer(in_channels: 3, out_channels: 32, kernel_size: 9, stride: 1)
+    public var in1 = InstanceNorm2d<Float>(featureCount: 32)
+    public var conv2 = ConvLayer(in_channels: 32, out_channels: 64, kernel_size: 3, stride: 2)
+    public var in2 = InstanceNorm2d<Float>(featureCount: 64)
+    public var conv3 = ConvLayer(in_channels: 64, out_channels: 128, kernel_size: 3, stride: 2)
+    public var in3 = InstanceNorm2d<Float>(featureCount: 128)
+
+    // residual layers
+    public var res1 = ResidualBlock(channels: 128)
+    public var res2 = ResidualBlock(channels: 128)
+    public var res3 = ResidualBlock(channels: 128)
+    public var res4 = ResidualBlock(channels: 128)
+    public var res5 = ResidualBlock(channels: 128)
+
+    // upsampling layers
+    public var deconv1 = UpsampleConvLayer(in_channels: 128, out_channels: 64, kernel_size: 3, stride: 1, upsample: 2.0)
+    public var in4 = InstanceNorm2d<Float>(featureCount: 64)
+    public var deconv2 = UpsampleConvLayer(in_channels: 64, out_channels: 32, kernel_size: 3, stride: 1, upsample: 2.0)
+    public var in5 = InstanceNorm2d<Float>(featureCount: 32)
+    public var deconv3 = UpsampleConvLayer(in_channels: 32, out_channels: 3, kernel_size: 9, stride: 1)
+
+    // activation
+    public var relu = ReLU<Float>()
+
+    public init() {}
+
+    @differentiable
+    public func callAsFunction(_ input: Input) -> Output {
+        let convolved1 = input.sequenced(through: conv1, in1, relu)
+        let convolved2 = convolved1.sequenced(through: conv2, in2, relu)
+        let convolved3 = convolved2.sequenced(through: conv3, in3, relu)
+        let residual = convolved3.sequenced(through: res1, res2, res3, res4, res5)
+        let upscaled1 = residual.sequenced(through: deconv1, in4, relu)
+        let upscaled2 = upscaled1.sequenced(through: deconv2, in5)
+        let upscaled3 = deconv3(upscaled2)
+        return upscaled3
+    }
+}
+
+/// Helper convolution layer with padding
+public struct ConvLayer: Layer {
+    public typealias Input = Tensor<Float>
+    public typealias Output = Tensor<Float>
+
+    public var reflection_pad: ReflectionPad2d<Float>
+    public var conv2d: Conv2D<Float>
+
+    public init(in_channels: Int, out_channels: Int, kernel_size: Int, stride: Int) {
+        reflection_pad = ReflectionPad2d<Float>(padding: Int(kernel_size / 2))
+        conv2d = Conv2D<Float>(filterShape: (kernel_size, kernel_size, in_channels, out_channels), strides: (stride, stride))
+    }
+
+    @differentiable
+    public func callAsFunction(_ input: Input) -> Output {
+        return input.sequenced(through: reflection_pad, conv2d)
+    }
+
+}
+
+/// Residual block
+public struct ResidualBlock: Layer {
+    public typealias Input = Tensor<Float>
+    public typealias Output = Tensor<Float>
+
+    public var conv1: ConvLayer
+    public var in1: InstanceNorm2d<Float>
+    public var conv2: ConvLayer
+    public var in2: InstanceNorm2d<Float>
+    public var relu = ReLU<Float>()
+
+    public init(channels: Int) {
+        conv1 = ConvLayer(in_channels: channels, out_channels: channels, kernel_size: 3, stride: 1)
+        in1 = InstanceNorm2d<Float>(featureCount: channels)
+        conv2 = ConvLayer(in_channels: channels, out_channels: channels, kernel_size: 3, stride: 1)
+        in2 = InstanceNorm2d<Float>(featureCount: channels)
+    }
+
+    @differentiable
+    public func callAsFunction(_ input: Input) -> Output {
+        return input + input.sequenced(
+                through:
+                conv1, in1, relu,
+                conv2, in2
+        )
+    }
+}
+
+/// Upscaling layer
+public struct UpsampleConvLayer: Layer {
+    public typealias Input = Tensor<Float>
+    public typealias Output = Tensor<Float>
+
+    @noDerivative public let upsample: Float
+    public var reflection_pad: ReflectionPad2d<Float>
+    public var conv2d: Conv2D<Float>
+
+    public init(in_channels: Int, out_channels: Int, kernel_size: Int, stride: Int, upsample: Float = 1.0) {
+        self.upsample = upsample
+        reflection_pad = ReflectionPad2d<Float>(padding: Int(kernel_size / 2))
+        conv2d = Conv2D<Float>(filterShape: (kernel_size, kernel_size, in_channels, out_channels), strides: (stride, stride))
+    }
+
+    @differentiable
+    public func callAsFunction(_ input: Input) -> Output {
+        let resizedInput = resizeNearestNeighbor(input, scale_factor: upsample)
+        return resizedInput.sequenced(through: reflection_pad, conv2d)
+    }
+}
+

FastStyleTransfer/README.md

@@ -0,0 +1,21 @@
+# Fast Style Transfer
+
+Based on the [PyTorch implementation](https://github.com/pytorch/examples/tree/master/fast_neural_style).
+The model should be trainable, but so far it's only tested for inference with pre-trained weights (included in `Demo/weights`). 
+
+## Example
+Run demo application to apply styles to jpeg images:
+```
+swift run FastStyleTranserDemo --weights=FastStyleTranser/Demo/weights/candy.npz --input=FastStyleTranser/Demo/examples/cat.jpg --output=candy_cat.jpg
+swift run FastStyleTranserDemo --weights=FastStyleTranser/Demo/weights/mosaic.npz --input=FastStyleTranser/Demo/examples/cat.jpg --output=mosaic_cat.jpg
+```
+
+<img src="Demo/examples/cat.jpg" height="240" width="240" align="left">
+<img src="Demo/examples/cat_candy.jpg" height="240" width="240" align="left">
+<img src="Demo/examples/cat_mosaic.jpg" height="240" width="240">
+
+## Requirements
+Requires Python and NumPy to load weights.
+
+## Jupyter Notebook
+Run [demo notebook](Demo/ColabDemo.ipynb) in [Colab](https://colab.research.google.com/github/vvmnnnkv/swift-models/blob/fast-style/FastStyleTransfer/Demo/ColabDemo.ipynb)!