Skip to content

Commit

Permalink
Add image adapters and tests (google-gemini#80)
Browse files Browse the repository at this point in the history
* Add image  adapters and tests

* code review feedback
  • Loading branch information
morganchen12 authored Jan 9, 2024
1 parent 3e2302a commit 7442ebf
Show file tree
Hide file tree
Showing 4 changed files with 116 additions and 5 deletions.
2 changes: 2 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -36,3 +36,5 @@ Mint
# CLI Tool
Apps/GoogleAICLI/GoogleAICLI.xcodeproj/xcshareddata/xcschemes/*
GenerativeAI-Info.plist

xcodebuild.log
4 changes: 0 additions & 4 deletions Package.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
44 changes: 43 additions & 1 deletion Sources/GoogleAI/PartsRepresentable.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand Down Expand Up @@ -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 []
}
Expand All @@ -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 []
}
}
71 changes: 71 additions & 0 deletions Tests/GoogleAITests/PartsRepresentableTests.swift
Original file line number Diff line number Diff line change
@@ -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
}

0 comments on commit 7442ebf

Please sign in to comment.