Replace TensorFlow image loading with stb_image (#429)

* Replace TensorFlow image loading with Swim.

* Formatting of Image.swift.

* Adding Swim to CMake dependencies.

* Adding tests for image loading and saving.

* Replacing Swim with stb_image.

* Removing Swim references in main CMakeLists.

* Converting grayscale color channels to match RGB expectations.

* Add directory prefix to CMake files.

* Adding C language support to CMake.

* Adding an StbImage modulemap.

* Move StbImage CMakeLists into its own directory, make import @_implementationOnly, update with Saleem's suggestions.

* Fix paths.

* Renamed to STBImage, moved the modulemap.

Author: BradLarson

.gitignore

@@ -5,7 +5,6 @@
 
 .build
 *.xcodeproj
-*.png
 .DS_Store
 .swiftpm
 cifar-10-batches-py/

CMakeLists.txt

@@ -1,6 +1,6 @@
 cmake_minimum_required(VERSION 3.16)
 project(Models
-  LANGUAGES Swift)
+  LANGUAGES C Swift)
 
 if(CMAKE_VERSION VERSION_LESS 3.17)
   if(NOT CMAKE_SYSTEM_NAME STREQUAL Windows)

Package.swift

@@ -42,7 +42,10 @@ let package = Package(
     targets: [
         .target(name: "Batcher", path: "Batcher"),
         .target(name: "Datasets", dependencies: ["ModelSupport", "Batcher"], path: "Datasets"),
-        .target(name: "ModelSupport", dependencies: ["SwiftProtobuf"], path: "Support"),
+        .target(name: "STBImage", path: "Support/STBImage"),
+        .target(
+            name: "ModelSupport", dependencies: ["SwiftProtobuf", "STBImage"], path: "Support",
+            exclude: ["STBImage"]),
         .target(name: "ImageClassificationModels", path: "Models/ImageClassification"),
         .target(name: "VideoClassificationModels", path: "Models/Spatiotemporal"),
         .target(name: "TextModels", dependencies: ["Datasets"], path: "Models/Text"),

Support/CMakeLists.txt

@@ -1,3 +1,5 @@
+add_subdirectory(STBImage)
+
 add_library(ModelSupport
   BijectiveDictionary.swift
   Checkpoints/CheckpointIndexReader.swift
@@ -22,6 +24,7 @@ set_target_properties(ModelSupport PROPERTIES
 target_compile_options(ModelSupport PRIVATE
   $<$<BOOL:${BUILD_TESTING}>:-enable-testing>)
 target_link_libraries(ModelSupport PUBLIC
+  STBImage
   SwiftProtobuf)
 
 

Support/Image.swift

@@ -13,14 +13,23 @@
 // limitations under the License.
 
 import Foundation
+@_implementationOnly import STBImage
 import TensorFlow
 
+// Image loading and saving is inspired by t-ae's Swim library: https://github.com/t-ae/swim
+// and uses the stb_image single-file C headers from https://github.com/nothings/stb .
+
 public struct Image {
     public enum ByteOrdering {
         case bgr
         case rgb
     }
 
+    public enum Colorspace {
+        case rgb
+        case grayscale
+    }
+
     enum ImageTensor {
         case float(data: Tensor<Float>)
         case uint8(data: Tensor<UInt8>)
@@ -44,20 +53,41 @@ public struct Image {
     }
 
     public init(jpeg url: URL, byteOrdering: ByteOrdering = .rgb) {
-        let loadedFile = _Raw.readFile(filename: StringTensor(url.absoluteString))
-        let loadedJpeg = _Raw.decodeJpeg(contents: loadedFile, channels: 3, dctMethod: "")
         if byteOrdering == .bgr {
-            self.imageData = .uint8(
-                data: _Raw.reverse(loadedJpeg, dims: Tensor<Bool>([false, false, false, true])))
+            // TODO: Add BGR byte reordering.
+            fatalError("BGR byte ordering is currently unsupported.")
         } else {
-            self.imageData = .uint8(data: loadedJpeg)
+            guard FileManager.default.fileExists(atPath: url.path) else {
+                // TODO: Proper error propagation for this.
+                fatalError("File does not exist at: \(url.path).")
+            }
+            
+            var width: Int32 = 0
+            var height: Int32 = 0
+            var bpp: Int32 = 0
+            guard let bytes = stbi_load(url.path, &width, &height, &bpp, 0) else {
+                // TODO: Proper error propagation for this.
+                fatalError("Unable to read image at: \(url.path).")
+            }
+
+            let data = [UInt8](UnsafeBufferPointer(start: bytes, count: Int(width * height * bpp)))
+            stbi_image_free(bytes)
+            var loadedTensor = Tensor<UInt8>(
+                shape: [Int(height), Int(width), Int(bpp)], scalars: data)
+            if bpp == 1 {
+                loadedTensor = loadedTensor.broadcasted(to: [Int(height), Int(width), 3])
+            }
+            self.imageData = .uint8(data: loadedTensor)
         }
     }
 
-    public func save(to url: URL, format: _Raw.Format = .rgb, quality: Int64 = 95) {
+    public func save(to url: URL, format: Colorspace = .rgb, quality: Int64 = 95) {
         let outputImageData: Tensor<UInt8>
+        let bpp: Int32
+
         switch format {
         case .grayscale:
+            bpp = 1
             switch self.imageData {
             case let .uint8(data): outputImageData = data
             case let .float(data):
@@ -67,51 +97,51 @@ public struct Image {
                 outputImageData = Tensor<UInt8>(adjustedData)
             }
         case .rgb:
+            bpp = 3
             switch self.imageData {
             case let .uint8(data): outputImageData = data
             case let .float(data):
-                outputImageData = Tensor<UInt8>(
-                    _Raw.clipByValue(t: data, clipValueMin: Tensor(0), clipValueMax: Tensor(255)))
+                outputImageData = Tensor<UInt8>(data.clipped(min: 0, max: 255))
+            }
+        }
+        
+        let height = Int32(outputImageData.shape[0])
+        let width = Int32(outputImageData.shape[1])
+        outputImageData.scalars.withUnsafeBufferPointer { bytes in
+            let status = stbi_write_jpg(
+                url.path, width, height, bpp, bytes.baseAddress!, Int32(quality))
+            guard status != 0 else {
+                // TODO: Proper error propagation for this.
+                fatalError("Unable to save image to: \(url.path).")
             }
-        default:
-            print("Image saving isn't supported for the format \(format).")
-            exit(-1)
         }
-
-        let encodedJpeg = _Raw.encodeJpeg(
-            image: outputImageData, format: format, quality: quality, xmpMetadata: "")
-        _Raw.writeFile(filename: StringTensor(url.absoluteString), contents: encodedJpeg)
     }
 
     public func resized(to size: (Int, Int)) -> Image {
         switch self.imageData {
         case let .uint8(data):
-            return Image(
-                tensor: _Raw.resizeBilinear(
-                    images: Tensor<UInt8>([data]),
-                    size: Tensor<Int32>([Int32(size.0), Int32(size.1)])).squeezingShape(at: 0))
+            let resizedImage = resize(images: Tensor<Float>(data), size: size, method: .bilinear)
+            return Image(tensor: Tensor<UInt8>(resizedImage))
         case let .float(data):
-            return Image(
-                tensor: _Raw.resizeBilinear(
-                    images: Tensor<Float>([data]),
-                    size: Tensor<Int32>([Int32(size.0), Int32(size.1)])).squeezingShape(at: 0))
+            let resizedImage = resize(images: data, size: size, method: .bilinear)
+            return Image(tensor: resizedImage)
         }
-
     }
 }
 
-public func saveImage(_ tensor: Tensor<Float>, shape: (Int, Int), size: (Int, Int)? = nil,
-                      format: _Raw.Format = .rgb, directory: String, name: String,
-                      quality: Int64 = 95) throws {
+public func saveImage(
+    _ tensor: Tensor<Float>, shape: (Int, Int), size: (Int, Int)? = nil,
+    format: Image.Colorspace = .rgb, directory: String, name: String,
+    quality: Int64 = 95
+) throws {
     try createDirectoryIfMissing(at: directory)
+
     let channels: Int
     switch format {
     case .rgb: channels = 3
     case .grayscale: channels = 1
-    default:
-        print("\(format) is not supported yet.")
-        exit(-1)
     }
+
     let reshapedTensor = tensor.reshaped(to: [shape.0, shape.1, channels])
     let image = Image(tensor: reshapedTensor)
     let resizedImage = size != nil ? image.resized(to: (size!.0, size!.1)) : image

Support/STBImage/CMakeLists.txt

@@ -0,0 +1,5 @@
+add_library(STBImage STATIC
+  stb_image_write.c
+  stb_image.c)
+target_include_directories(STBImage PUBLIC
+  include)

Support/STBImage/include/module.modulemap

@@ -0,0 +1,5 @@
+module STBImage {
+    header "stb_image.h"
+    header "stb_image_write.h"
+    export *
+}