Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add image adapters and tests #80

Merged
merged 2 commits into from
Jan 9, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
}
Loading