Add JPEG loading / saving via Raw TF operations, removing Matplotlib dependencies

Autoencoder/main.swift

@@ -59,12 +59,16 @@ for epoch in 1...epochCount {
         shape: [1, imageHeight * imageWidth], scalars: dataset.trainingImages[epoch].scalars)
     let testImage = autoencoder(sampleImage)
 
-    saveImage(
-        tensor: sampleImage, size: (imageWidth, imageHeight), directory: outputFolder,
-        name: "epoch-\(epoch)-input")
-    saveImage(
-        tensor: testImage, size: (imageWidth, imageHeight), directory: outputFolder,
-        name: "epoch-\(epoch)-output")
+    do {
+        try saveImage(
+            sampleImage, size: (imageWidth, imageHeight), directory: outputFolder,
+            name: "epoch-\(epoch)-input")
+        try saveImage(
+            testImage, size: (imageWidth, imageHeight), directory: outputFolder,
+            name: "epoch-\(epoch)-output")
+    } catch {
+        print("Could not save image with error: \(error)")
+    }
 
     let sampleLoss = meanSquaredError(predicted: testImage, expected: sampleImage)
     print("[Epoch: \(epoch)] Loss: \(sampleLoss)")

GAN/main.swift

@@ -117,11 +117,11 @@ let optD = Adam(for: discriminator, learningRate: 2e-4, beta1: 0.5)
 let testImageGridSize = 4
 let testVector = sampleVector(size: testImageGridSize * testImageGridSize)
 
-func saveImageGrid(_ testImage: Tensor<Float>, name: String) {
+func saveImageGrid(_ testImage: Tensor<Float>, name: String) throws {
     var gridImage = testImage.reshaped(
         to: [
             testImageGridSize, testImageGridSize,
-            imageHeight, imageWidth
+            imageHeight, imageWidth,
         ])
     // Add padding.
     gridImage = gridImage.padded(forSizes: [(0, 0), (0, 0), (1, 1), (1, 1)], with: 1)
@@ -130,13 +130,13 @@ func saveImageGrid(_ testImage: Tensor<Float>, name: String) {
     gridImage = gridImage.reshaped(
         to: [
             (imageHeight + 2) * testImageGridSize,
-            (imageWidth + 2) * testImageGridSize
+            (imageWidth + 2) * testImageGridSize,
         ])
     // Convert [-1, 1] range to [0, 1] range.
     gridImage = (gridImage + 1) / 2
 
-    saveImage(
-        tensor: gridImage, size: (gridImage.shape[0], gridImage.shape[1]), directory: outputFolder,
+    try saveImage(
+        gridImage, size: (gridImage.shape[0], gridImage.shape[1]), directory: outputFolder,
         name: name)
 }
 
@@ -176,7 +176,12 @@ for epoch in 1...epochCount {
     // Start inference phase.
     Context.local.learningPhase = .inference
     let testImage = generator(testVector)
-    saveImageGrid(testImage, name: "epoch-\(epoch)-output")
+
+    do {
+        try saveImageGrid(testImage, name: "epoch-\(epoch)-output")
+    } catch {
+        print("Could not save image grid with error: \(error)")
+    }
 
     let lossG = generatorLoss(fakeLogits: testImage)
     print("[Epoch: \(epoch)] Loss-G: \(lossG)")

Support/FileManagement.swift

@@ -14,11 +14,10 @@
 
 import Foundation
 
-public func createDirectoryIfMissing(path: String) {
-    if !FileManager.default.fileExists(atPath: path) {
-        try! FileManager.default.createDirectory(
-            atPath: path,
-            withIntermediateDirectories: false,
-            attributes: nil)
-    }
+public func createDirectoryIfMissing(path: String) throws {
+    guard !FileManager.default.fileExists(atPath: path) else { return }
+    try FileManager.default.createDirectory(
+        atPath: path,
+        withIntermediateDirectories: false,
+        attributes: nil)
 }

Support/Image.swift

@@ -21,12 +21,12 @@ public struct Image {
         case rgb
     }
 
-    enum ImageTensorFormat {
+    enum ImageTensor {
         case float(data: Tensor<Float>)
         case uint8(data: Tensor<UInt8>)
     }
 
-    let imageData: ImageTensorFormat
+    let imageData: ImageTensor
 
     public init(tensor: Tensor<UInt8>) {
         self.imageData = .uint8(data: tensor)
@@ -48,7 +48,7 @@ public struct Image {
     }
 
     public func save(to url: URL, quality: Int64 = 95) {
-        // Currently only saving in grayscale 
+        // This currently only saves in grayscale.
         let outputImageData: Tensor<UInt8>
         switch self.imageData {
         case let .uint8(data): outputImageData = data
@@ -81,9 +81,8 @@ public struct Image {
     }
 }
 
-public func saveImage(tensor: Tensor<Float>, size: (Int, Int), directory: String, name: String) {
-    createDirectoryIfMissing(path: directory)
-
+public func saveImage(_ tensor: Tensor<Float>, size: (Int, Int), directory: String, name: String) throws {
+    try createDirectoryIfMissing(path: directory)
     let reshapedTensor = tensor.reshaped(to: [size.0, size.1, 1])
     let image = Image(tensor: reshapedTensor)
     let outputURL = URL(fileURLWithPath:"\(directory)\(name).jpg")