Add image adapters and tests (#80)

* Add image  adapters and tests

* code review feedback

Author: morganchen12

.gitignore

@@ -36,3 +36,5 @@ Mint
 # CLI Tool
 Apps/GoogleAICLI/GoogleAICLI.xcodeproj/xcshareddata/xcschemes/*
 GenerativeAI-Info.plist
+
+xcodebuild.log

Package.swift

@@ -25,16 +25,12 @@ let package = Package(
     .macCatalyst(.v15),
   ],
   products: [
-    // Products define the executables and libraries a package produces, making them visible to
-    // other packages.
     .library(
       name: "GoogleGenerativeAI",
       targets: ["GoogleGenerativeAI"]
     ),
   ],
   targets: [
-    // Targets are the basic building blocks of a package, defining a module or a test suite.
-    // Targets can depend on other targets in this package and products from dependencies.
     .target(
       name: "GoogleGenerativeAI",
       path: "Sources"

Sources/GoogleAI/PartsRepresentable.swift

@@ -13,12 +13,15 @@
 // limitations under the License.
 
 import Foundation
+import UniformTypeIdentifiers
 #if canImport(UIKit)
   import UIKit // For UIImage extensions.
 #elseif canImport(AppKit)
   import AppKit // For NSImage extensions.
 #endif
 
+private let imageCompressionQuality: CGFloat = 0.8
+
 /// A protocol describing any data that could be interpreted as model input data.
 public protocol PartsRepresentable {
   var partsValue: [ModelContent.Part] { get }
@@ -50,7 +53,7 @@ extension [any PartsRepresentable]: PartsRepresentable {
   /// Enables images to be representable as ``PartsRepresentable``.
   extension UIImage: PartsRepresentable {
     public var partsValue: [ModelContent.Part] {
-      guard let data = jpegData(compressionQuality: 0.8) else {
+      guard let data = jpegData(compressionQuality: imageCompressionQuality) else {
         Logging.default.error("[GoogleGenerativeAI] Couldn't create JPEG from UIImage.")
         return []
       }
@@ -77,3 +80,42 @@ extension [any PartsRepresentable]: PartsRepresentable {
     }
   }
 #endif
+
+extension CGImage: PartsRepresentable {
+  public var partsValue: [ModelContent.Part] {
+    let output = NSMutableData()
+    guard let imageDestination = CGImageDestinationCreateWithData(
+      output, UTType.jpeg.identifier as CFString, 1, nil
+    ) else {
+      Logging.default.error("[GoogleGenerativeAI] Couldn't create JPEG from CGImage.")
+      return []
+    }
+    CGImageDestinationAddImage(imageDestination, self, nil)
+    CGImageDestinationSetProperties(imageDestination, [
+      kCGImageDestinationLossyCompressionQuality: imageCompressionQuality,
+    ] as CFDictionary)
+    if CGImageDestinationFinalize(imageDestination) {
+      return [.data(mimetype: "image/jpeg", output as Data)]
+    }
+    Logging.default.error("[GoogleGenerativeAI] Couldn't create JPEG from CGImage.")
+    return []
+  }
+}
+
+extension CIImage: PartsRepresentable {
+  public var partsValue: [ModelContent.Part] {
+    let context = CIContext()
+    let jpegData = (colorSpace ?? CGColorSpace(name: CGColorSpace.sRGB))
+      .flatMap {
+        // The docs specify kCGImageDestinationLossyCompressionQuality as a supported option, but
+        // Swift's type system does not allow this.
+        // [kCGImageDestinationLossyCompressionQuality: imageCompressionQuality]
+        context.jpegRepresentation(of: self, colorSpace: $0, options: [:])
+      }
+    if let jpegData = jpegData {
+      return [.data(mimetype: "image/jpeg", jpegData)]
+    }
+    Logging.default.error("[GoogleGenerativeAI] Couldn't create JPEG from CIImage.")
+    return []
+  }
+}

Tests/GoogleAITests/PartsRepresentableTests.swift

@@ -0,0 +1,71 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import CoreGraphics
+import CoreImage
+import XCTest
+#if canImport(UIKit)
+  import UIKit
+#else
+  import AppKit
+#endif
+
+final class PartsRepresentableTests: XCTestCase {
+  func testModelContentFromCGImageIsNotEmpty() throws {
+    // adapted from https://forums.swift.org/t/creating-a-cgimage-from-color-array/18634/2
+    var srgbArray = [UInt32](repeating: 0xFFFF_FFFF, count: 8 * 8)
+    let image = srgbArray.withUnsafeMutableBytes { ptr -> CGImage in
+      let ctx = CGContext(
+        data: ptr.baseAddress,
+        width: 8,
+        height: 8,
+        bitsPerComponent: 8,
+        bytesPerRow: 4 * 8,
+        space: CGColorSpace(name: CGColorSpace.sRGB)!,
+        bitmapInfo: CGBitmapInfo.byteOrder32Little.rawValue +
+          CGImageAlphaInfo.premultipliedFirst.rawValue
+      )!
+      return ctx.makeImage()!
+    }
+    let modelContent = image.partsValue
+    XCTAssert(modelContent.count > 0, "Expected non-empty model content for CGImage: \(image)")
+  }
+
+  func testModelContentFromCIImageIsNotEmpty() throws {
+    let image = CIImage(color: CIColor.red)
+      .cropped(to: CGRect(origin: CGPointZero, size: CGSize(width: 16, height: 16)))
+    let modelContent = image.partsValue
+    XCTAssert(modelContent.count > 0, "Expected non-empty model content for CGImage: \(image)")
+  }
+
+  #if canImport(UIKit)
+    func testModelContentFromUIImageIsNotEmpty() throws {
+      let coreImage = CIImage(color: CIColor.red)
+        .cropped(to: CGRect(origin: CGPointZero, size: CGSize(width: 16, height: 16)))
+      let image = UIImage(ciImage: coreImage)
+      let modelContent = image.partsValue
+      XCTAssert(modelContent.count > 0, "Expected non-empty model content for UIImage: \(image)")
+    }
+  #else
+    func testModelContentFromNSImageIsNotEmpty() throws {
+      let coreImage = CIImage(color: CIColor.red)
+        .cropped(to: CGRect(origin: CGPointZero, size: CGSize(width: 16, height: 16)))
+      let rep = NSCIImageRep(ciImage: coreImage)
+      let image = NSImage(size: rep.size)
+      image.addRepresentation(rep)
+      let modelContent = image.partsValue
+      XCTAssert(modelContent.count > 0, "Expected non-empty model content for NSImage: \(image)")
+    }
+  #endif
+}