diff --git a/Package.swift b/Package.swift index 1b032aa..9fc5e4f 100644 --- a/Package.swift +++ b/Package.swift @@ -50,6 +50,9 @@ let package = Package( .executableTarget(name: "SortDecisionsCheck", dependencies: ["SortAnything"]), .executableTarget(name: "SortDecisionsDemo", dependencies: ["SortAnything"], exclude: ["README.md"]), .executableTarget(name: "SortAnythingDemo", dependencies: ["SortAnything"], exclude: ["README.md"]), + .target(name: "ImageSort", dependencies: ["FluidUse"]), + .executableTarget(name: "ImageSortCheck", dependencies: ["ImageSort", "FluidUse"]), + .executableTarget(name: "ImageSortDemo", dependencies: ["ImageSort"], exclude: ["README.md"]), .testTarget( name: "FluidUseTests", dependencies: ["FluidUse", "LayaTetris"], resources: [.copy("Fixtures")] diff --git a/Sources/FluidUse/SigLIP2/SigLIP2ImagePreprocessor.swift b/Sources/FluidUse/SigLIP2/SigLIP2ImagePreprocessor.swift new file mode 100644 index 0000000..95123cd --- /dev/null +++ b/Sources/FluidUse/SigLIP2/SigLIP2ImagePreprocessor.swift @@ -0,0 +1,95 @@ +import CoreGraphics +import Foundation + +/// Matches the Hugging Face SigLIP image processor: PIL bilinear resize (antialiased when shrinking) to a square, +/// rescale to [0, 1], normalize per channel. Output is planar float32 `[3, size, size]`. +public enum SigLIP2ImagePreprocessor { + public static func pixels(from image: CGImage, config: SigLIP2Config) throws -> [Float] { + let width = image.width + let height = image.height + let size = config.imageSize + var rgba = [UInt8](repeating: 0, count: width * height * 4) + guard + let space = CGColorSpace(name: CGColorSpace.sRGB), + let context = CGContext( + data: &rgba, width: width, height: height, bitsPerComponent: 8, bytesPerRow: width * 4, + space: space, bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue) + else { + throw SigLIP2Error.invalidInput("Could not decode a \(width)×\(height) image") + } + context.draw(image, in: CGRect(x: 0, y: 0, width: width, height: height)) + // Horizontal pass, then vertical, each rounded to 8 bits like PIL. Layout stays 4 bytes per pixel. + let horizontal = resample( + rgba, lines: height, inLength: width, outLength: size, sampleStride: 4, lineStride: width * 4, + outSampleStride: 4, outLineStride: size * 4, outCount: height * size * 4) + let resized = resample( + horizontal, lines: size, inLength: height, outLength: size, sampleStride: size * 4, lineStride: 4, + outSampleStride: size * 4, outLineStride: 4, outCount: size * size * 4) + var planar = [Float](repeating: 0, count: 3 * size * size) + for channel in 0..<3 { + let scale = 1 / (255 * config.imageStd[channel]) + let offset = config.imageMean[channel] / config.imageStd[channel] + let base = channel * size * size + for pixel in 0..<(size * size) { + planar[base + pixel] = Float(resized[pixel * 4 + channel]) * scale - offset + } + } + return planar + } + + /// One separable pass of PIL's `ImagingResample` with the triangle filter, rounding to 8 bits like PIL. + static func resample( + _ input: [UInt8], lines: Int, inLength: Int, outLength: Int, sampleStride: Int, lineStride: Int, + outSampleStride: Int, outLineStride: Int, outCount: Int + ) -> [UInt8] { + let scale = Double(inLength) / Double(outLength) + let filterScale = max(scale, 1) + var starts = [Int](repeating: 0, count: outLength) + var counts = [Int](repeating: 0, count: outLength) + let taps = Int((filterScale * 2).rounded(.up)) + 2 + var weights = [Float](repeating: 0, count: outLength * taps) + for out in 0.. 0 ? weight / total : 0) + } + starts[out] = low + counts[out] = min(high - low, taps) + } + var output = [UInt8](repeating: 0, count: outCount) + input.withUnsafeBufferPointer { source in + output.withUnsafeMutableBufferPointer { destination in + weights.withUnsafeBufferPointer { weight in + for line in 0.. SigLIP2Manager { + let directory = try await SigLIP2ModelStore.ensure(cacheDirectory: cacheDirectory, progress: progress) + return try await load(from: directory, computeUnits: computeUnits) + } + + /// Loads the manager from `SIGLIP2_MODEL_DIR` when set, otherwise from the published packages. + public static func loadDefault(progress: SigLIP2ModelStore.Progress? = nil) async throws -> SigLIP2Manager { + if let path = ProcessInfo.processInfo.environment["SIGLIP2_MODEL_DIR"], !path.isEmpty { + return try await load(from: URL(fileURLWithPath: path)) + } + return try await load(progress: progress) + } + + /// Loads `config.json`, `tokenizer.json`, and the image and text packages (`.mlmodelc` preferred) from + /// `directory`, as written by the mobius converter. + public static func load( + from directory: URL, computeUnits: MLComputeUnits = .cpuAndNeuralEngine + ) async throws + -> SigLIP2Manager + { + let configURL = directory.appendingPathComponent("config.json") + guard let configData = try? Data(contentsOf: configURL) else { + throw SigLIP2Error.invalidAsset("Missing config.json in \(directory.path)") + } + let config = try JSONDecoder().decode(SigLIP2Config.self, from: configData) + let tokenizer = try SigLIP2Tokenizer( + tokenizerJsonURL: directory.appendingPathComponent("tokenizer.json"), length: config.textLength) + let configuration = MLModelConfiguration() + configuration.computeUnits = computeUnits + async let image = loadModel(named: "\(config.name)-image-\(config.precision)", in: directory, configuration) + async let text = loadModel(named: "\(config.name)-text-\(config.precision)", in: directory, configuration) + return try await SigLIP2Manager(config: config, tokenizer: tokenizer, imageModel: image, textModel: text) + } + + private static func loadModel( + named name: String, in directory: URL, _ configuration: MLModelConfiguration + ) async throws -> MLModel { + let compiled = directory.appendingPathComponent("\(name).mlmodelc") + let package = directory.appendingPathComponent("\(name).mlpackage") + let url: URL + if FileManager.default.fileExists(atPath: compiled.path) { + url = compiled + } else if FileManager.default.fileExists(atPath: package.path) { + url = try await MLModel.compileModel(at: package) + } else { + throw SigLIP2Error.invalidAsset("Missing \(name).mlmodelc or .mlpackage in \(directory.path)") + } + return try await MLModel.load(contentsOf: url, configuration: configuration) + } + + /// L2-normalized text embedding per label. Compute once per label set and reuse. + public func embed(labels: [String]) async throws -> [[Float]] { + var embeddings: [[Float]] = [] + for label in labels { + let ids = try tokenizer.encode(label) + let input = try MLMultiArray(shape: [1, NSNumber(value: ids.count)], dataType: .int32) + let pointer = input.dataPointer.assumingMemoryBound(to: Int32.self) + for (index, id) in ids.enumerated() { pointer[index] = id } + let output = try await textModel.prediction( + from: MLDictionaryFeatureProvider(dictionary: ["input_ids": MLFeatureValue(multiArray: input)])) + embeddings.append(try Self.vector(output, name: "text_embeds")) + } + return embeddings + } + + /// Image embedding plus when the Core ML call began and ended (`DispatchTime` uptime nanoseconds). + public struct TimedEmbedding: Sendable { + public let embedding: [Float] + public let predictionStart: UInt64 + public let predictionEnd: UInt64 + } + + /// L2-normalized image embedding. + public func embed(image: CGImage) async throws -> [Float] { try await embedTimed(image: image).embedding } + + /// L2-normalized image embedding, with the timing of the model call alone (no decoding or resizing). + public func embedTimed(image: CGImage) async throws -> TimedEmbedding { + let pixels = try SigLIP2ImagePreprocessor.pixels(from: image, config: config) + let size = NSNumber(value: config.imageSize) + let input = try MLMultiArray(shape: [1, 3, size, size], dataType: .float32) + pixels.withUnsafeBufferPointer { source in + input.dataPointer.assumingMemoryBound(to: Float.self).update(from: source.baseAddress!, count: pixels.count) + } + let features = try MLDictionaryFeatureProvider(dictionary: ["pixel_values": MLFeatureValue(multiArray: input)]) + let start = DispatchTime.now().uptimeNanoseconds + let output = try await imageModel.prediction(from: features) + let end = DispatchTime.now().uptimeNanoseconds + return TimedEmbedding( + embedding: try Self.vector(output, name: "image_embeds"), predictionStart: start, predictionEnd: end) + } + + /// Scores `image` against label embeddings from `embed(labels:)`. + public func classify(image: CGImage, labels: [String], labelEmbeddings: [[Float]]) async throws -> SigLIP2Answer { + guard labels.count == labelEmbeddings.count, !labels.isEmpty else { + throw SigLIP2Error.invalidInput("Expected one embedding per label") + } + return score(imageEmbedding: try await embed(image: image), labels: labels, labelEmbeddings: labelEmbeddings) + } + + public func score(imageEmbedding: [Float], labels: [String], labelEmbeddings: [[Float]]) -> SigLIP2Answer { + let similarities = labelEmbeddings.map { label in zip(label, imageEmbedding).reduce(0) { $0 + $1.0 * $1.1 } } + let probabilities = similarities.map { 1 / (1 + exp(-(config.logitScale * $0 + config.logitBias))) } + return SigLIP2Answer(labels: labels, similarities: similarities, probabilities: probabilities) + } + + private static func vector(_ output: MLFeatureProvider, name: String) throws -> [Float] { + guard let array = output.featureValue(for: name)?.multiArrayValue else { + throw SigLIP2Error.predictionFailed("Missing \(name)") + } + return (0.. Void + + public static let repository = "FluidInference/siglip2-base-patch16-256-coreml" + static let revision = "524a5a7d666f23002853831915cf1cc13734f73f" + + private struct Asset { + let path: String + let sha256: String + } + + private static let assets = [ + Asset(path: "config.json", sha256: "b5d7aaa84399aaa277f9cc00c6c401edc05d535f47f4feb1880e0e4ed2a1a8d4"), + Asset( + path: "siglip2-base-patch16-256-image-fp16.mlpackage/Data/com.apple.CoreML/model.mlmodel", + sha256: "702fdc0b557984e16e4bbcce4eec3568f08a7459bd3c27c770a487823bda808b"), + Asset( + path: "siglip2-base-patch16-256-image-fp16.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + sha256: "da086438b60ada3566f8c91bd8a632462f28d344186e7f60cb5abd223d2e99a8"), + Asset( + path: "siglip2-base-patch16-256-image-fp16.mlpackage/Manifest.json", + sha256: "964a40aa63c2d201a30f0aeab68aef8abc95c854a34f4cae533051cd99996e72"), + Asset( + path: "siglip2-base-patch16-256-text-fp16.mlpackage/Data/com.apple.CoreML/model.mlmodel", + sha256: "69fce0f538fbe78c45b953e1cc396fb912d701c6f985ac825bdec92833aca452"), + Asset( + path: "siglip2-base-patch16-256-text-fp16.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + sha256: "9a52dd8222973b6b4cdf55983689bb2fa18a27fcd20b83f8d8fb89de544530b5"), + Asset( + path: "siglip2-base-patch16-256-text-fp16.mlpackage/Manifest.json", + sha256: "a5dfacf23259261d32c5af26c5cfe4c41a44f0584739b2ef10e70c7f11c2a620"), + Asset( + path: "tokenizer_config.json", sha256: "9c8a03337138d3b5509e4c032f6863e769b7448750718372c907407d67f6a91b"), + Asset(path: "tokenizer.json", sha256: "caefd63119539a63be2d55ef3e05023fbb793948c4bda5bc0c366b42a382f903"), + ] + + /// Ensures the packages, tokenizer, and config exist and match their checksums; returns their directory. + public static func ensure(cacheDirectory: URL? = nil, progress: Progress? = nil) async throws -> URL { + let root = cacheDirectory ?? LayaModelStore.defaultCacheDirectory() + let directory = root.appendingPathComponent("siglip2-base-patch16-256-coreml") + let manager = FileManager.default + try manager.createDirectory(at: directory, withIntermediateDirectories: true) + for asset in assets { + let destination = directory.appendingPathComponent(asset.path) + if manager.fileExists(atPath: destination.path), try checksum(of: destination) == asset.sha256 { + continue + } + try manager.createDirectory(at: destination.deletingLastPathComponent(), withIntermediateDirectories: true) + progress?(asset.path, 0) + let escaped = asset.path.addingPercentEncoding(withAllowedCharacters: .urlPathAllowed) ?? asset.path + guard let url = URL(string: "https://huggingface.co/\(repository)/resolve/\(revision)/\(escaped)") else { + throw SigLIP2Error.invalidAsset("Invalid Hugging Face asset URL") + } + let (temporary, response) = try await URLSession.shared.download(from: url) + defer { try? manager.removeItem(at: temporary) } + guard let http = response as? HTTPURLResponse, http.statusCode == 200 else { + throw SigLIP2Error.invalidAsset("Download failed for \(asset.path)") + } + let actual = try checksum(of: temporary) + guard actual == asset.sha256 else { + throw SigLIP2Error.invalidAsset( + "Checksum mismatch for \(asset.path): expected \(asset.sha256), got \(actual)") + } + let size = (try manager.attributesOfItem(atPath: temporary.path)[.size] as? NSNumber)?.int64Value ?? 0 + try LayaModelStore.installDownloadedFile(temporary, at: destination) + progress?(asset.path, size) + } + return directory + } + + private static func checksum(of file: URL) throws -> String { + let handle = try FileHandle(forReadingFrom: file) + defer { try? handle.close() } + var digest = SHA256() + while let chunk = try handle.read(upToCount: 1_048_576), !chunk.isEmpty { + digest.update(data: chunk) + } + return digest.finalize().map { String(format: "%02x", $0) }.joined() + } +} diff --git a/Sources/FluidUse/SigLIP2/SigLIP2Tokenizer.swift b/Sources/FluidUse/SigLIP2/SigLIP2Tokenizer.swift new file mode 100644 index 0000000..7f0ae26 --- /dev/null +++ b/Sources/FluidUse/SigLIP2/SigLIP2Tokenizer.swift @@ -0,0 +1,82 @@ +import Foundation + +/// Gemma BPE tokenizer as used by SigLIP 2's text encoder: spaces become `▁`, ranked merges over Unicode scalars, +/// UTF-8 byte fallback for scalars outside the vocabulary, `` appended, `` to a fixed length. +/// Text is lowercased first, matching SigLIP 2's training. +public final class SigLIP2Tokenizer: Sendable { + public let length: Int + public let padId: Int + public let eosId: Int + + private let vocab: [String: Int] + private let ranks: [String: Int] + + public convenience init(tokenizerJsonURL: URL, length: Int = 64) throws { + try self.init(data: Data(contentsOf: tokenizerJsonURL), length: length) + } + + public init(data: Data, length: Int = 64) throws { + guard + let root = try JSONSerialization.jsonObject(with: data) as? [String: Any], + let model = root["model"] as? [String: Any], + let vocab = model["vocab"] as? [String: Int], + let merges = model["merges"] as? [[String]] + else { + throw SigLIP2Error.invalidAsset("tokenizer.json is not a Gemma BPE tokenizer") + } + guard let padId = vocab[""], let eosId = vocab[""] else { + throw SigLIP2Error.invalidAsset("tokenizer.json has no or token") + } + var ranks: [String: Int] = [:] + ranks.reserveCapacity(merges.count) + for (rank, pair) in merges.enumerated() where pair.count == 2 { + ranks[Self.key(pair[0], pair[1])] = rank + } + self.vocab = vocab + self.ranks = ranks + self.length = length + self.padId = padId + self.eosId = eosId + } + + /// Token ids for `text`, ``-terminated and padded to `length`. + public func encode(_ text: String) throws -> [Int32] { + var ids = tokenize(text.lowercased()) + ids.append(eosId) + guard ids.count <= length else { + throw SigLIP2Error.invalidInput("Label needs \(ids.count) tokens; the text encoder takes \(length)") + } + ids.append(contentsOf: repeatElement(padId, count: length - ids.count)) + return ids.map(Int32.init) + } + + /// Token ids without `` or padding. + public func tokenize(_ text: String) -> [Int] { + let normalized = text.replacingOccurrences(of: " ", with: "\u{2581}") + var symbols = normalized.unicodeScalars.map { String($0) } + while symbols.count > 1 { + var best: (rank: Int, index: Int)? + for index in 0..<(symbols.count - 1) { + if let rank = ranks[Self.key(symbols[index], symbols[index + 1])], rank < (best?.rank ?? .max) { + best = (rank, index) + } + } + guard let best else { break } + symbols[best.index] += symbols[best.index + 1] + symbols.remove(at: best.index + 1) + } + var ids: [Int] = [] + for symbol in symbols { + if let id = vocab[symbol] { + ids.append(id) + } else { + for byte in symbol.utf8 { + ids.append(vocab[String(format: "<0x%02X>", byte)] ?? vocab[""] ?? 3) + } + } + } + return ids + } + + private static func key(_ left: String, _ right: String) -> String { left + "\u{0}" + right } +} diff --git a/Sources/FluidUse/SigLIP2/SigLIP2Types.swift b/Sources/FluidUse/SigLIP2/SigLIP2Types.swift new file mode 100644 index 0000000..1bb6986 --- /dev/null +++ b/Sources/FluidUse/SigLIP2/SigLIP2Types.swift @@ -0,0 +1,55 @@ +import Foundation + +public enum SigLIP2Error: Error, LocalizedError { + case invalidAsset(String) + case invalidInput(String) + case predictionFailed(String) + + public var errorDescription: String? { + switch self { + case .invalidAsset(let reason): "Invalid SigLIP 2 asset: \(reason)" + case .invalidInput(let reason): "Invalid SigLIP 2 input: \(reason)" + case .predictionFailed(let reason): "SigLIP 2 prediction failed: \(reason)" + } + } +} + +/// Preprocessing and scoring constants written by the converter (`config.json`). +public struct SigLIP2Config: Codable, Sendable { + public let modelId: String + public let imageSize: Int + public let imageMean: [Float] + public let imageStd: [Float] + public let textLength: Int + public let logitScale: Float + public let logitBias: Float + public let embeddingDim: Int + public let precision: String + + enum CodingKeys: String, CodingKey { + case modelId = "model_id" + case imageSize = "image_size" + case imageMean = "image_mean" + case imageStd = "image_std" + case textLength = "text_length" + case logitScale = "logit_scale" + case logitBias = "logit_bias" + case embeddingDim = "embedding_dim" + case precision + } + + /// Package base name, e.g. `siglip2-base-patch16-256`. + public var name: String { modelId.split(separator: "/").last.map(String.init) ?? modelId } +} + +/// One image scored against a label set. +public struct SigLIP2Answer: Sendable { + public let labels: [String] + /// Cosine similarity per label. + public let similarities: [Float] + /// Independent `sigmoid(scale · cos + bias)` per label, as SigLIP was trained. + public let probabilities: [Float] + + public var selectedIndex: Int { similarities.indices.max { similarities[$0] < similarities[$1] } ?? 0 } + public var selectedLabel: String { labels[selectedIndex] } +} diff --git a/Sources/ImageSort/ImageSorter.swift b/Sources/ImageSort/ImageSorter.swift new file mode 100644 index 0000000..14e986d --- /dev/null +++ b/Sources/ImageSort/ImageSorter.swift @@ -0,0 +1,71 @@ +import CoreGraphics +import FluidUse +import Foundation +import ImageIO + +/// Sorts photos into breeds with SigLIP 2 on Core ML. Calls are not serialized: several `sort` calls may be in flight. +public final class ImageSorter: Sendable { + public struct Result: Sendable { + public let breed: String + /// SigLIP's own sigmoid probability for the chosen label. + public let probability: Float + /// Softmax of the scaled similarities across all labels: the chosen label's share among the candidates. + public let share: Float + /// The five most likely breeds with their shares, best first. + public let top: [(breed: String, share: Float)] + public let milliseconds: Double + /// When the Core ML call itself began and ended (`DispatchTime` uptime nanoseconds). + public let predictionStart: UInt64 + public let predictionEnd: UInt64 + } + + public let breeds: [String] + private let manager: SigLIP2Manager + private let embeddings: [[Float]] + + private init(manager: SigLIP2Manager, breeds: [String], embeddings: [[Float]]) { + self.manager = manager + self.breeds = breeds + self.embeddings = embeddings + } + + /// Loads the encoders (from `SIGLIP2_MODEL_DIR`, or downloaded once from Hugging Face) and embeds the breed + /// prompts once. + public static func load( + breeds: [String] = PetsSample.breeds, progress: SigLIP2ModelStore.Progress? = nil + ) async throws -> ImageSorter { + let manager = try await SigLIP2Manager.loadDefault(progress: progress) + let embeddings = try await manager.embed(labels: breeds.map(PetsSample.prompt(for:))) + return ImageSorter(manager: manager, breeds: breeds, embeddings: embeddings) + } + + public var modelName: String { manager.config.name } + + public func sort(_ item: PetItem) async throws -> Result { + let image = try Self.decode(item.file) + let start = DispatchTime.now().uptimeNanoseconds + let timed = try await manager.embedTimed(image: image) + let answer = manager.score(imageEmbedding: timed.embedding, labels: breeds, labelEmbeddings: embeddings) + let milliseconds = Double(DispatchTime.now().uptimeNanoseconds - start) / 1e6 + let scale = manager.config.logitScale + let best = answer.similarities[answer.selectedIndex] + let weights = answer.similarities.map { exp(scale * ($0 - best)) } + let total = weights.reduce(0, +) + let ranked = weights.indices.sorted { weights[$0] > weights[$1] }.prefix(5).map { + (breed: breeds[$0], share: weights[$0] / total) + } + return Result( + breed: answer.selectedLabel, probability: answer.probabilities[answer.selectedIndex], share: 1 / total, + top: ranked, milliseconds: milliseconds, predictionStart: timed.predictionStart, + predictionEnd: timed.predictionEnd) + } + + public static func decode(_ file: URL) throws -> CGImage { + guard let source = CGImageSourceCreateWithURL(file as CFURL, nil), + let image = CGImageSourceCreateImageAtIndex(source, 0, nil) + else { + throw SigLIP2Error.invalidInput("Could not decode \(file.lastPathComponent)") + } + return image + } +} diff --git a/Sources/ImageSort/PetsSample.swift b/Sources/ImageSort/PetsSample.swift new file mode 100644 index 0000000..1dadcfb --- /dev/null +++ b/Sources/ImageSort/PetsSample.swift @@ -0,0 +1,169 @@ +import Foundation + +/// One photo from the Oxford-IIIT Pets test split with its gold breed. +public struct PetItem: Codable, Sendable, Identifiable, Hashable { + public let id: Int + public let breed: String + /// Cached JPEG on disk. + public let file: URL +} + +/// Seeded sample of the Oxford-IIIT Pets test split (CC BY-SA 4.0), fetched from the Hugging Face dataset viewer +/// API on first use and cached locally. Nothing is bundled. +public enum PetsSample { + public static let attribution = "Oxford-IIIT Pets (Parkhi et al., 2012) · CC BY-SA 4.0" + public static let testCount = 3669 + + /// The 37 breeds in dataset label order. + public static let breeds = [ + "abyssinian", "american bulldog", "american pit bull terrier", "basset hound", "beagle", "bengal", "birman", + "bombay", "boxer", "british shorthair", "chihuahua", "egyptian mau", "english cocker spaniel", + "english setter", "german shorthaired", "great pyrenees", "havanese", "japanese chin", "keeshond", + "leonberger", "maine coon", "miniature pinscher", "newfoundland", "persian", "pomeranian", "pug", "ragdoll", + "russian blue", "saint bernard", "samoyed", "scottish terrier", "shiba inu", "siamese", "sphynx", + "staffordshire bull terrier", "wheaten terrier", "yorkshire terrier", + ] + + /// Prompt per breed, as scored in the mobius Pets check. + public static func prompt(for breed: String) -> String { "a photo of a \(breed), a type of pet." } + + static let endpoint = "https://datasets-server.huggingface.co/rows" + static let pageSize = 100 + + private struct Page: Decodable { + struct Entry: Decodable { + struct Row: Decodable { + struct Image: Decodable { let src: String } + let image: Image + let label: Int + } + let rowIndex: Int + let row: Row + + enum CodingKeys: String, CodingKey { + case row + case rowIndex = "row_idx" + } + } + let rows: [Entry] + } + + public static func cacheDirectory() -> URL { + FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask)[0] + .appendingPathComponent("FluidUse/image-sort/oxford-pets-test") + } + + /// `count` photos in a seeded shuffled order (every eligible photo when `count` is nil); `testOnly` restricts + /// the pool to the 3,669 test photos. The viewer API fetches the + /// 3,669 test photos; a cache that also holds the train split (ids from 3,669) samples from both. + public static func load( + count: Int? = 1000, seed: UInt64 = 0, testOnly: Bool = false, + progress: (@Sendable (Int, Int) -> Void)? = nil + ) async throws -> [PetItem] { + let directory = cacheDirectory() + let manifestURL = directory.appendingPathComponent("manifest.json") + let manager = FileManager.default + try manager.createDirectory(at: directory, withIntermediateDirectories: true) + + var labels: [Int: Int] = [:] + if let data = try? Data(contentsOf: manifestURL), + let saved = try? JSONDecoder().decode([Int: Int].self, from: data), + saved.count >= testCount + { + labels = saved + } + var sources: [Int: String] = [:] + if labels.count < testCount { + for offset in stride(from: 0, to: testCount, by: pageSize) { + for entry in try await page(offset: offset) { + labels[entry.rowIndex] = entry.row.label + sources[entry.rowIndex] = entry.row.image.src + } + } + try JSONEncoder().encode(labels).write(to: manifestURL) + } + + var generator = SeededGenerator(seed: seed) + let pool = labels.keys.filter { !testOnly || $0 < testCount }.sorted() + let chosen = Array(pool.shuffled(using: &generator).prefix(count ?? pool.count)) + let missing = chosen.filter { !manager.fileExists(atPath: file(for: $0).path) } + if !missing.isEmpty { + if sources.isEmpty { + for offset in stride(from: 0, to: testCount, by: pageSize) { + for entry in try await page(offset: offset) { sources[entry.rowIndex] = entry.row.image.src } + } + } + try await download( + missing, sources: sources, done: chosen.count - missing.count, total: chosen.count, progress) + } + return chosen.map { PetItem(id: $0, breed: breeds[labels[$0]!], file: file(for: $0)) } + } + + static func file(for row: Int) -> URL { cacheDirectory().appendingPathComponent("\(row).jpg") } + + private static func page(offset: Int) async throws -> [Page.Entry] { + var components = URLComponents(string: endpoint)! + components.queryItems = [ + URLQueryItem(name: "dataset", value: "timm/oxford-iiit-pet"), + URLQueryItem(name: "config", value: "default"), + URLQueryItem(name: "split", value: "test"), + URLQueryItem(name: "offset", value: String(offset)), + URLQueryItem(name: "length", value: String(pageSize)), + ] + let data = try await fetch(components.url!) + return try JSONDecoder().decode(Page.self, from: data).rows + } + + /// GET with backoff: the dataset viewer answers bursts with 429 and a `Retry-After`. + static func fetch(_ url: URL, attempts: Int = 8) async throws -> Data { + var delay = 2.0 + for attempt in 1...attempts { + let (data, response) = try await URLSession.shared.data(from: url) + let http = response as? HTTPURLResponse + if http?.statusCode == 200 { return data } + guard attempt < attempts, http?.statusCode == 429 || (http?.statusCode ?? 0) >= 500 else { break } + let wait = (http?.value(forHTTPHeaderField: "Retry-After")).flatMap(Double.init) ?? delay + try await Task.sleep(for: .seconds(min(wait, 60))) + delay *= 2 + } + throw URLError(.badServerResponse) + } + + private static func download( + _ rows: [Int], sources: [Int: String], done: Int, total: Int, _ progress: (@Sendable (Int, Int) -> Void)? + ) async throws { + try await withThrowingTaskGroup(of: Void.self) { group in + var next = 0 + var finished = done + func launch() { + guard next < rows.count else { return } + let row = rows[next] + next += 1 + group.addTask { + guard let source = sources[row], let url = URL(string: source) else { throw URLError(.badURL) } + try await fetch(url).write(to: file(for: row), options: .atomic) + } + } + for _ in 0..<8 { launch() } + while try await group.next() != nil { + finished += 1 + progress?(finished, total) + launch() + } + } + } +} + +/// xorshift64*, so a seed gives the same sample on every machine. +struct SeededGenerator: RandomNumberGenerator { + private var state: UInt64 + + init(seed: UInt64) { state = seed &+ 0x9E37_79B9_7F4A_7C15 } + + mutating func next() -> UInt64 { + state ^= state >> 12 + state ^= state << 25 + state ^= state >> 27 + return state &* 0x2545_F491_4F6C_DD1D + } +} diff --git a/Sources/ImageSortCheck/main.swift b/Sources/ImageSortCheck/main.swift new file mode 100644 index 0000000..8d61ae7 --- /dev/null +++ b/Sources/ImageSortCheck/main.swift @@ -0,0 +1,77 @@ +import FluidUse +import Foundation +import ImageSort + +// Headless checks for the SigLIP 2 image sorter. +// ImageSortCheck tokenizer token ids must equal the Python tokenizer's +// ImageSortCheck [--split=test|all] [--count=N] [--inflight=N] [--predictions=out.json] +// zero-shot Pets accuracy and speed; test split (3,669) by default, all 7,349 when the train split is cached +setvbuf(stdout, nil, _IOLBF, 0) +let arguments = CommandLine.arguments.dropFirst() + +func option(_ name: String) -> String? { + arguments.first { $0.hasPrefix("--\(name)=") }.map { String($0.dropFirst(name.count + 3)) } +} + +if arguments.first == "tokenizer", let path = arguments.dropFirst().first { + struct Case: Decodable { + let text: String + let ids: [Int32] + } + guard let directory = ProcessInfo.processInfo.environment["SIGLIP2_MODEL_DIR"] else { + fatalError("Set SIGLIP2_MODEL_DIR") + } + let tokenizer = try SigLIP2Tokenizer( + tokenizerJsonURL: URL(fileURLWithPath: directory).appendingPathComponent("tokenizer.json")) + let cases = try JSONDecoder().decode([Case].self, from: Data(contentsOf: URL(fileURLWithPath: path))) + var mismatches = 0 + for item in cases { + let ids = try tokenizer.encode(item.text) + if ids != item.ids { + mismatches += 1 + if mismatches <= 5 { + print("mismatch: \(item.text.debugDescription)\n swift \(ids)\n python \(item.ids)") + } + } + } + print("tokenizer: \(cases.count - mismatches)/\(cases.count) identical") + exit(mismatches == 0 ? 0 : 1) +} + +let count = option("count").flatMap(Int.init) +let inFlight = option("inflight").flatMap(Int.init) ?? 4 +let testOnly = option("split") != "all" +let items = try await PetsSample.load(count: count, testOnly: testOnly) { done, total in + if done % 250 == 0 || done == total { print("cached \(done)/\(total) photos") } +} +let sorter = try await ImageSorter.load() +_ = try await sorter.sort(items[0]) + +let start = Date() +let results = try await withThrowingTaskGroup(of: (PetItem, ImageSorter.Result).self) { group in + var next = 0 + var results: [(PetItem, ImageSorter.Result)] = [] + func launch() { + guard next < items.count else { return } + let item = items[next] + next += 1 + group.addTask { (item, try await sorter.sort(item)) } + } + for _ in 0.. [String: CGRect]) { + value.merge(nextValue()) { $1 } + } +} + +extension View { + fileprivate func reportFrame(_ key: String) -> some View { + background( + GeometryReader { proxy in + Color.clear.preference(key: FramesKey.self, value: [key: proxy.frame(in: .named("board"))]) + }) + } +} + +struct ContentView: View { + @EnvironmentObject private var model: ImageSortModel + @State private var frames: [String: CGRect] = [:] + + var body: some View { + GeometryReader { proxy in + // Narrow windows merge the header and the current photo into one band so the chart gets the rest. + let narrow = proxy.size.width < 980 + VStack(spacing: 0) { + if narrow { compactTop } else { header } + Divider() + switch model.phase { + case .loading(let message): + status(message, spinning: true) + case .failed(let message): + status("Failed: \(message)", spinning: false) + default: + board(narrow: narrow) + } + Divider() + footer + } + .coordinateSpace(name: "board") + .onPreferenceChange(FramesKey.self) { frames = $0 } + .overlay(alignment: .topLeading) { flying } + } + .background(Color(nsColor: .windowBackgroundColor)) + } + + private var compactTop: some View { + HStack(alignment: .top, spacing: 16) { + photo.frame(width: 176, height: 176) + VStack(alignment: .leading, spacing: 12) { + HStack(alignment: .top) { + titleBlock + Spacer(minLength: 8) + controls + } + stats(size: 18) + topFive.frame(maxWidth: 420, alignment: .leading) + } + } + .padding(.horizontal, 16) + .padding(.vertical, 12) + } + + private var header: some View { + ViewThatFits(in: .horizontal) { + HStack(alignment: .center, spacing: 18) { + titleBlock + Spacer(minLength: 12) + stats() + controls + } + VStack(alignment: .leading, spacing: 10) { + titleBlock + HStack(alignment: .center, spacing: 14) { + stats() + Spacer(minLength: 8) + controls + } + } + } + .padding(.horizontal, 20) + .padding(.vertical, 12) + } + + private var titleBlock: some View { + VStack(alignment: .leading, spacing: 2) { + Text("Sort photos").font(.system(size: 26, weight: .bold)) + Text("SigLIP 2 · Core ML on the Neural Engine · 37 breeds, zero-shot") + .font(.callout).foregroundStyle(.secondary).lineLimit(1).fixedSize() + } + } + + private func stats(size: CGFloat = 20) -> some View { + HStack(spacing: 0) { + stat("Sorted", "\(model.sorted)", size: size) + stat("Elapsed time", String(format: "%.1f s", model.elapsed), size: size) + stat("Photos/s", String(format: "%.0f", model.sorted > 0 ? model.photosPerSecond : 0), size: size) + stat("ms/photo", millisecondsPerPhoto, size: size) + stat("Accuracy", String(format: "%.1f%%", (model.accuracy ?? 0) * 100), size: size) + } + } + + /// Wall time per photo with several model calls overlapping. + private var millisecondsPerPhoto: String { + model.sorted > 0 ? String(format: "%.1f", 1000 / model.photosPerSecond) : "0.0" + } + + /// Equal-width tile, label over value, both centered, so columns stay put as numbers change. + private func stat(_ title: String, _ value: String, size: CGFloat) -> some View { + VStack(spacing: 3) { + Text(title.uppercased()).font(.caption2.weight(.semibold)).foregroundStyle(.secondary) + .lineLimit(1).fixedSize() + Text(value).font(.system(size: size, weight: .semibold, design: .rounded)).monospacedDigit() + .contentTransition(.numericText()).lineLimit(1).fixedSize() + } + .frame(width: size < 20 ? 104 : 120) + } + + private var controls: some View { + VStack(alignment: .leading, spacing: 6) { + HStack { + Button(model.phase == .running ? "Pause" : "Start") { model.toggleRun() } + .keyboardShortcut(.space, modifiers: []) + .disabled(!(model.phase == .ready || model.phase == .running || model.phase == .paused)) + .buttonStyle(.borderedProminent) + Button("Reset") { model.reset() }.disabled(model.phase == .running) + } + } + } + + private func board(narrow: Bool) -> some View { + Group { + if narrow { + chart(labelWidth: 150) + } else { + HStack(alignment: .top, spacing: 18) { + nowSorting.frame(width: 250) + chart(labelWidth: 190) + } + } + } + .padding(14) + } + + /// Photos travelling from the Now Sorting panel to their tile; they shrink to tile size on arrival. + private var flying: some View { + ZStack(alignment: .topLeading) { + if let from = frames["photo"], let chart = frames["chart"] { + let scale = chart.width / CGFloat(model.chartGeometry.width) + ForEach(model.flights) { flight in + let target = CGRect( + x: chart.minX + flight.slot.minX * scale, y: chart.minY + flight.slot.minY * scale, + width: flight.slot.width * scale, height: flight.slot.height * scale) + let side = flight.arrived ? max(target.width, 4) : from.width * 0.55 + Color.clear + .frame(width: side, height: side) + .overlay { + if let image = flight.image { + Image(decorative: image, scale: 1).resizable().aspectRatio(contentMode: .fill) + } + } + .clipShape(RoundedRectangle(cornerRadius: flight.arrived ? 1 : 8)) + .overlay( + RoundedRectangle(cornerRadius: flight.arrived ? 1 : 8) + .stroke(flight.wrong ? Color.red : Color.green, lineWidth: flight.arrived ? 1 : 3) + ) + .shadow(color: .black.opacity(0.4), radius: flight.arrived ? 0 : 6) + .position( + x: flight.arrived ? target.midX : from.midX, y: flight.arrived ? target.midY : from.midY) + } + } + } + .allowsHitTesting(false) + } + + private var photo: some View { + Color.secondary.opacity(0.08) + .aspectRatio(1, contentMode: .fit) + .overlay { + if let image = model.currentImage { + Image(decorative: image, scale: 1).resizable().aspectRatio(contentMode: .fill) + } + } + .clipShape(RoundedRectangle(cornerRadius: 12)) + .reportFrame("photo") + .overlay( + RoundedRectangle(cornerRadius: 12) + .stroke(model.current.map { $0.matchesGold ? Color.green : Color.red } ?? .clear, lineWidth: 3)) + } + + @ViewBuilder + private var topFive: some View { + if let current = model.current { + VStack(alignment: .leading, spacing: 5) { + ForEach(Array(current.result.top.enumerated()), id: \.offset) { rank, entry in + topRow(entry.breed, share: entry.share, first: rank == 0, gold: current.item.breed) + } + if !current.matchesGold { + Text("label: \(current.item.breed)").font(.caption.weight(.semibold)).foregroundStyle(.red) + } + } + } + } + + private var nowSorting: some View { + VStack(alignment: .leading, spacing: 10) { + Text("NOW SORTING · \(model.remaining) left").font(.caption.weight(.semibold)).foregroundStyle(.secondary) + photo + topFive + Spacer(minLength: 0) + Text("Labels are the 37 breed names, typed once: \"a photo of a {breed}, a type of pet.\"") + .font(.caption).foregroundStyle(.secondary) + } + } + + private func topRow(_ breed: String, share: Float, first: Bool, gold: String) -> some View { + HStack(spacing: 6) { + Text(breed).font(first ? .callout.weight(.bold) : .caption).lineLimit(1) + .frame(width: 128, alignment: .leading) + GeometryReader { proxy in + Capsule().fill(breed == gold ? Color.green : color(for: breed)) + .frame(width: max(2, proxy.size.width * CGFloat(share))) + } + .frame(height: first ? 10 : 6) + Text("\(Int(share * 100))%").font(.caption).monospacedDigit().frame(width: 34, alignment: .trailing) + } + } + + private func chart(labelWidth: CGFloat) -> some View { + GeometryReader { proxy in + let geometry = model.chartGeometry + let width = CGFloat(geometry.width) + let height = CGFloat(geometry.height) + let available = CGSize(width: proxy.size.width - labelWidth, height: proxy.size.height) + let scale = min(available.width / width, available.height / height) + let rowHeight = CGFloat(geometry.rowHeight) * scale + HStack(alignment: .top, spacing: 0) { + VStack(alignment: .trailing, spacing: 0) { + ForEach(model.breeds, id: \.self) { breed in + HStack(spacing: 6) { + Text(breed).lineLimit(1).minimumScaleFactor(0.7).foregroundStyle(color(for: breed)) + Text("\(model.counts[breed] ?? 0)").monospacedDigit().foregroundStyle(.secondary) + .frame(width: 30, alignment: .trailing) + } + .font(.system(size: max(8, min(13, rowHeight * 0.6)), weight: .semibold)) + .frame(width: labelWidth - 8, height: rowHeight, alignment: .trailing) + .padding(.trailing, 8) + } + } + ZStack(alignment: .topLeading) { + if let image = model.chartImage { + Image(decorative: image, scale: 1).resizable().interpolation(.medium) + .frame(width: width * scale, height: height * scale) + } + } + .frame(width: width * scale, height: height * scale, alignment: .topLeading) + .reportFrame("chart") + } + .onAppear { model.fitChart(to: available) } + .onChange(of: available) { _, size in model.fitChart(to: size) } + } + } + + private func status(_ message: String, spinning: Bool) -> some View { + VStack(spacing: 12) { + if spinning { ProgressView() } + Text(message).foregroundStyle(.secondary) + } + .frame(maxWidth: .infinity, maxHeight: .infinity) + } + + private var footer: some View { + HStack { + Text(PetsSample.attribution) + Spacer() + Text("Model: google/siglip2-base-patch16-256 (Apache-2.0), converted to Core ML by FluidInference") + } + .font(.caption).foregroundStyle(.secondary) + .padding(.horizontal, 20).padding(.vertical, 8) + } + + private func color(for breed: String) -> Color { ImageSortModel.cats.contains(breed) ? .orange : .blue } +} diff --git a/Sources/ImageSortDemo/ImageSortDemoApp.swift b/Sources/ImageSortDemo/ImageSortDemoApp.swift new file mode 100644 index 0000000..a8bb5e2 --- /dev/null +++ b/Sources/ImageSortDemo/ImageSortDemoApp.swift @@ -0,0 +1,28 @@ +import AppKit +import SwiftUI + +@main +struct ImageSortDemoApp: App { + @StateObject private var model = ImageSortModel() + + init() { + setvbuf(stdout, nil, _IOLBF, 0) + // Bare SwiftPM executables start as background processes; make this one a regular windowed app. + for key in UserDefaults.standard.dictionaryRepresentation().keys where key.hasPrefix("NSWindow Frame") { + UserDefaults.standard.removeObject(forKey: key) + } + NSApplication.shared.setActivationPolicy(.regular) + NSApplication.shared.activate(ignoringOtherApps: true) + } + + var body: some Scene { + WindowGroup("Sort photos — SigLIP 2 on-device") { + ContentView() + .environmentObject(model) + .frame(minWidth: 640, minHeight: 560) + .task { await model.prepare() } + } + .defaultSize(width: 1560, height: 980) + .windowResizability(.contentMinSize) + } +} diff --git a/Sources/ImageSortDemo/ImageSortModel.swift b/Sources/ImageSortDemo/ImageSortModel.swift new file mode 100644 index 0000000..6989bd9 --- /dev/null +++ b/Sources/ImageSortDemo/ImageSortModel.swift @@ -0,0 +1,338 @@ +import CoreGraphics +import Foundation +import ImageIO +import ImageSort +import SwiftUI + +/// Drives the stream: loads Pets photos, sorts them with SigLIP 2, and grows the photo chart. +@MainActor +final class ImageSortModel: ObservableObject { + enum Phase: Equatable { + case loading(String) + case ready + case running + case paused + case finished + case failed(String) + } + + struct Placed: Identifiable, Sendable { + let item: PetItem + let result: ImageSorter.Result + let tile: CGImage? + var id: Int { item.id } + var matchesGold: Bool { result.breed == item.breed } + } + + @Published private(set) var phase: Phase = .loading("Starting…") + @Published private(set) var remaining = 0 + @Published private(set) var sorted = 0 + @Published private(set) var correct = 0 + @Published private(set) var elapsed: Double = 0 + /// Chart slot of the photo shown in the Now Sorting panel. + private var currentSlot: CGRect? + @Published private(set) var counts: [String: Int] = [:] + @Published private(set) var chartImage: CGImage? + /// The photo just sorted, shown large with its top-5 breeds. + @Published private(set) var current: Placed? + @Published private(set) var currentImage: CGImage? + /// Photos travelling from the Now Sorting panel to their chart slot. + @Published private(set) var flights: [Flight] = [] + + struct Flight: Identifiable { + let id: Int + let image: CGImage? + /// Target tile in chart pixels. + let slot: CGRect + let wrong: Bool + var arrived = false + } + + static let travel = 0.4 + static let turboFlightsPerFlush = 2 + private var flightCounter = 0 + + let breeds = PetsSample.breeds + let total: Int + static let turboInFlight = 4 + static let turboFlush = 0.05 + static let cats: Set = [ + "abyssinian", "bengal", "birman", "bombay", "british shorthair", "egyptian mau", "maine coon", "persian", + "ragdoll", "russian blue", "siamese", "sphynx", + ] + + private let environment = ProcessInfo.processInfo.environment + /// IMAGE_SORT_LOG=1 prints every decision to stdout (for a terminal next to the window). + private let logDecisions = ProcessInfo.processInfo.environment["IMAGE_SORT_LOG"] == "1" + private var chart: PhotoChart? + private var queue: [PetItem] = [] + private var landed: Set = [] + private var items: [PetItem] = [] + private var sorter: ImageSorter? + private var runner: Task? + private var modelMilliseconds: [Double] = [] + private var runStart: Date? + private var elapsedBeforePause: Double = 0 + private var preparing = false + + init() { + total = ProcessInfo.processInfo.environment["IMAGE_SORT_COUNT"].flatMap(Int.init) ?? 1000 + } + + var photosPerSecond: Double { elapsed > 0 ? Double(sorted) / elapsed : 0 } + var accuracy: Double? { sorted > 0 ? Double(correct) / Double(sorted) : nil } + var medianMilliseconds: Double? { + guard !modelMilliseconds.isEmpty else { return nil } + return modelMilliseconds.sorted()[modelMilliseconds.count / 2] + } + + func prepare() async { + guard !preparing else { return } + preparing = true + do { + phase = .loading("Fetching \(total) Oxford-IIIT Pets photos…") + items = try await PetsSample.load(count: total) { [weak self] done, all in + Task { @MainActor in self?.phase = .loading("Cached \(done) / \(all) photos") } + } + phase = .loading("Loading SigLIP 2 (first run downloads 750 MB) and embedding 37 breed names…") + sorter = try await ImageSorter.load { [weak self] file, bytes in + guard bytes > 0 else { return } + Task { @MainActor in self?.phase = .loading("Downloaded \(file)") } + } + _ = try await sorter?.sort(items[0]) + reset() + print("ready at \(Date().timeIntervalSince1970)") + // IMAGE_SORT_WAIT=1 keeps the chart empty until Start (Space). + if environment["IMAGE_SORT_WAIT"] != "1" { + try? await Task.sleep(for: .seconds(1.5)) + toggleRun() + } + } catch { + phase = .failed(error.localizedDescription) + } + } + + /// Pixel size and shape of the chart bitmap, for the view's layout. + struct ChartGeometry: Equatable { + var width = 1 + var height = 1 + var rowHeight = 1 + } + + @Published private(set) var chartGeometry = ChartGeometry() + private var levels = 2 + + private func makeChart() { + let chart = PhotoChart(rows: breeds.count, levels: levels) { [breeds] row in + Self.cats.contains(breeds[row]) + ? CGColor(red: 1, green: 0.6, blue: 0.2, alpha: 0.10) + : CGColor(red: 0.3, green: 0.55, blue: 1, alpha: 0.10) + } + self.chart = chart + chartGeometry = ChartGeometry(width: chart.width, height: chart.height, rowHeight: chart.rowHeight) + chartImage = chart.snapshot() + } + + /// Reshapes the empty chart to fill `size` (points available for the tiles). Ignored once sorting has begun. + func fitChart(to size: CGSize) { + let best = PhotoChart.bestLevels(rows: breeds.count, size: size) + guard best != levels, sorted == 0, flights.isEmpty, phase != .running, chart != nil else { return } + levels = best + makeChart() + } + + func reset() { + runner?.cancel() + runner = nil + makeChart() + queue = items + remaining = items.count + landed = [] + counts = [:] + current = nil + currentImage = nil + currentSlot = nil + flights = [] + sorted = 0 + correct = 0 + elapsed = 0 + elapsedBeforePause = 0 + modelMilliseconds = [] + phase = .ready + } + + func toggleRun() { + switch phase { + case .running: + runner?.cancel() + runner = nil + elapsedBeforePause = elapsed + phase = .paused + case .ready, .paused: + phase = .running + runStart = Date() + runner = Task { [weak self] in await self?.run() } + default: + break + } + } + + private func run() async { + guard let sorter else { return } + while !Task.isCancelled, !queue.isEmpty { + await runTurbo(sorter) + } + if !Task.isCancelled, queue.isEmpty { + tick() + phase = .finished + // The last photo leaves the panel for its tile too, leaving only the finished chart. + if let current, let slot = currentSlot { + fly(image: currentImage ?? current.tile, to: slot, wrong: !current.matchesGold) + } + current = nil + currentImage = nil + currentSlot = nil + print( + "finished \(sorted) photos in \(String(format: "%.2f", elapsed)) s " + + "(\(String(format: "%.0f", photosPerSecond)) photos/s), " + + "correct \(String(format: "%.1f", (accuracy ?? 0) * 100))%") + } + } + + private func runTurbo(_ sorter: ImageSorter) async { + var pending: [Placed] = [] + var lastFlush = Date() + for await placed in Self.stream(sorter, items: queue, inFlight: Self.turboInFlight) { + pending.append(placed) + if Date().timeIntervalSince(lastFlush) >= Self.turboFlush { + flush(pending) + pending.removeAll(keepingCapacity: true) + lastFlush = Date() + } + if Task.isCancelled { break } + } + flush(pending) + } + + /// Lands a batch at once; a couple of its photos also fly so the motion stays visible. + private func flush(_ batch: [Placed]) { + let slots = land(batch) + for (placed, slot) in slots.suffix(Self.turboFlightsPerFlush) where flights.count < 24 { + fly(image: placed.tile, to: slot, wrong: !placed.matchesGold) + } + } + + /// Animates `image` from the Now Sorting panel to `slot`; the tile is already drawn there. + private func fly( + image: CGImage?, to slot: CGRect, wrong: Bool + ) { + flightCounter += 1 + let id = flightCounter + flights.append(Flight(id: id, image: image, slot: slot, wrong: wrong)) + let travel = Self.travel + Task { @MainActor [weak self] in + try? await Task.sleep(for: .milliseconds(20)) + withAnimation(.easeIn(duration: travel)) { + if let index = self?.flights.firstIndex(where: { $0.id == id }) { self?.flights[index].arrived = true } + } + try? await Task.sleep(for: .seconds(travel)) + self?.flights.removeAll { $0.id == id } + } + } + + /// Sorts `items` with `inFlight` calls always running off the main thread and streams each result. + private nonisolated static func stream( + _ sorter: ImageSorter, items: [PetItem], inFlight: Int + ) + -> AsyncStream + { + AsyncStream { continuation in + let producer = Task.detached { + await withTaskGroup(of: Placed?.self) { group in + var next = 0 + func launch() { + guard next < items.count, !Task.isCancelled else { return } + let item = items[next] + next += 1 + group.addTask { + guard let result = try? await sorter.sort(item) else { return nil } + return Placed(item: item, result: result, tile: thumbnail(item.file, size: 160)) + } + } + for _ in 0.. CGImage? { + guard let source = CGImageSourceCreateWithURL(file as CFURL, nil) else { return nil } + return CGImageSourceCreateThumbnailAtIndex( + source, 0, + [ + kCGImageSourceCreateThumbnailFromImageAlways: true, + kCGImageSourceThumbnailMaxPixelSize: size, + kCGImageSourceCreateThumbnailWithTransform: true, + ] as CFDictionary) + } + + /// Draws the batch into the chart and updates the statistics; returns each photo's slot. + @discardableResult + private func land(_ incoming: [Placed]) -> [(Placed, CGRect)] { + let batch = incoming.filter { self.landed.insert($0.id).inserted } + guard let last = batch.last, let chart else { return [] } + var updated = counts + var slots: [(Placed, CGRect)] = [] + for placed in batch { + let index = updated[placed.result.breed, default: 0] + let row = breeds.firstIndex(of: placed.result.breed) ?? 0 + chart.draw(placed.tile, row: row, index: index, wrong: !placed.matchesGold) + if let slot = chart.slot(row: row, index: index) { slots.append((placed, slot)) } + updated[placed.result.breed] = index + 1 + correct += placed.matchesGold ? 1 : 0 + modelMilliseconds.append(placed.result.milliseconds) + } + counts = updated + chartImage = chart.snapshot() + current = last + currentImage = Self.thumbnail(last.item.file, size: 480) + currentSlot = slots.last { $0.0.id == last.id }?.1 + let ids = Set(batch.map(\.id)) + queue.removeAll { ids.contains($0.id) } + remaining = queue.count + sorted += batch.count + tick() + if logDecisions { log(batch) } + return slots + } + + private func log(_ batch: [Placed]) { + let (cyan, yellow, red, green, dim, reset) = + ("\u{1B}[1;36m", "\u{1B}[33m", "\u{1B}[31m", "\u{1B}[32m", "\u{1B}[2m", "\u{1B}[0m") + // One photo per update keeps the terminal readable (and cheap) at hundreds of photos per second. + var lines = "" + for (offset, placed) in batch.enumerated().suffix(1) { + let number = sorted - batch.count + offset + 1 + let mark = placed.matchesGold ? "\(green)✓\(reset)" : "\(red)✗ label: \(placed.item.breed)\(reset)" + lines += "\(cyan)▶ #\(number) photo \(placed.item.id).jpg\(reset)\n" + lines += + " \(yellow)→ \(placed.result.breed)\(reset) · \(Int(placed.result.share * 100))% of 37 · " + + "\(red)model call \(String(format: "%.1f", placed.result.milliseconds)) ms\(reset) \(mark)\n" + } + let wrong = batch.filter { !$0.matchesGold }.count + lines += + "\(dim) +\(batch.count - 1) more (\(wrong) wrong) · sorted \(sorted)/\(total) · " + + "\(String(format: "%.1f", elapsed)) s\(reset)\n" + print(lines, terminator: "") + } + + private func tick() { + if let runStart { elapsed = elapsedBeforePause + Date().timeIntervalSince(runStart) } + } +} diff --git a/Sources/ImageSortDemo/PhotoChart.swift b/Sources/ImageSortDemo/PhotoChart.swift new file mode 100644 index 0000000..0864219 --- /dev/null +++ b/Sources/ImageSortDemo/PhotoChart.swift @@ -0,0 +1,86 @@ +import CoreGraphics +import Foundation + +/// A bar chart made of photos: one row per breed, each sorted photo drawn as a tile at the end of its row. +/// Tiles go into one bitmap, so thousands of them cost one image on screen. `levels` tiles stack in each row; +/// the column count follows so every row holds `capacity` photos. +@MainActor +final class PhotoChart { + static let tile = 16 + static let rowGap = 6 + static let capacity = 260 + + let rows: Int + let levels: Int + let columns: Int + let width: Int + let height: Int + var rowHeight: Int { Self.tile * levels + Self.rowGap } + private let context: CGContext + + init(rows: Int, levels: Int, tint: (Int) -> CGColor) { + self.rows = rows + self.levels = levels + columns = (Self.capacity + levels - 1) / levels + width = columns * Self.tile + height = rows * (Self.tile * levels + Self.rowGap) + context = CGContext( + data: nil, width: width, height: height, bitsPerComponent: 8, bytesPerRow: 0, + space: CGColorSpace(name: CGColorSpace.sRGB)!, bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue)! + for row in 0.. Int { + guard size.width > 0, size.height > 0 else { return 2 } + return (1...10).max { scale(rows: rows, levels: $0, size: size) < scale(rows: rows, levels: $1, size: size) } + ?? 2 + } + + private static func scale(rows: Int, levels: Int, size: CGSize) -> CGFloat { + let width = CGFloat((capacity + levels - 1) / levels * tile) + let height = CGFloat(rows * (tile * levels + rowGap)) + return min(size.width / width, size.height / height) + } + + /// Tile rectangle for the `index`-th photo of `row`, in top-left pixel coordinates (nil past capacity). + func slot(row: Int, index: Int) -> CGRect? { + guard index < Self.capacity else { return nil } + let column = index / levels + let level = index % levels + return CGRect( + x: column * Self.tile, y: row * rowHeight + level * Self.tile, width: Self.tile, height: Self.tile) + } + + /// Row a slot belongs to. + func row(of slot: CGRect) -> Int { Int(slot.minY) / rowHeight } + + func draw(_ image: CGImage?, row: Int, index: Int, wrong: Bool) { + guard let slot = slot(row: row, index: index) else { return } + let target = rect(top: Int(slot.minY), x: Int(slot.minX), width: Self.tile, height: Self.tile) + if let image { + context.draw(Self.squareCrop(image), in: target.insetBy(dx: 0.5, dy: 0.5)) + } + if wrong { + context.setStrokeColor(CGColor(red: 1, green: 0.2, blue: 0.2, alpha: 1)) + context.setLineWidth(2.5) + context.stroke(target.insetBy(dx: 1.25, dy: 1.25)) + } + } + + func snapshot() -> CGImage? { context.makeImage() } + + /// Converts a top-left rectangle to the context's bottom-left coordinates. + private func rect(top: Int, x: Int, width: Int, height: Int) -> CGRect { + CGRect(x: x, y: self.height - top - height, width: width, height: height) + } + + private static func squareCrop(_ image: CGImage) -> CGImage { + let side = min(image.width, image.height) + let crop = CGRect(x: (image.width - side) / 2, y: (image.height - side) / 2, width: side, height: side) + return image.cropping(to: crop) ?? image + } +} diff --git a/Sources/ImageSortDemo/README.md b/Sources/ImageSortDemo/README.md new file mode 100644 index 0000000..fcf5e49 --- /dev/null +++ b/Sources/ImageSortDemo/README.md @@ -0,0 +1,26 @@ +# Sort photos + +Sorts Oxford-IIIT Pets photos into 37 breeds with SigLIP 2 (base, 256 px) on Core ML. The only hint the model +gets is each breed's name, in the prompt `a photo of a {breed}, a type of pet.`; nothing is trained on these photos. + +```bash +swift run -c release ImageSortDemo +``` + +The first run downloads the Core ML packages (about 750 MB, checksum-verified) from +[FluidInference/siglip2-base-patch16-256-coreml](https://huggingface.co/FluidInference/siglip2-base-patch16-256-coreml) +and 1,000 test photos from the Hugging Face dataset viewer; both are cached. `SIGLIP2_MODEL_DIR` loads a local +conversion instead. + +- Each sorted photo flies from the current-photo panel into its breed's row, so the rows grow into a bar chart made + of photos; a red frame marks a photo whose true breed differs. Orange rows are cat breeds, blue are dog breeds. + The panel shows the latest photo with its five most likely breeds. +- Four photos are in flight at once, so decoding and resizing on the CPU overlap the encoder on the Neural Engine. +- `IMAGE_SORT_COUNT` sets the sample size (default 1,000; up to 3,669 test photos, or 7,349 when the train split is + also cached). `IMAGE_SORT_WAIT=1` waits for Start (Space); `IMAGE_SORT_LOG=1` prints decisions to the terminal. + +Headless: `swift run -c release ImageSortCheck` (M5 Pro, macOS 27: 94.85% on the 3,669 test photos, 202 photos/s; +`--split=all` 94.26% on 7,349). The same 7,349 photos with the PyTorch model (transformers, fp32, MPS, batch 32) +take 102 s at 72 photos/s and 4.2 GB peak memory, against 36 s, 202 photos/s, and 262 MB here. + +Data: Oxford-IIIT Pets (Parkhi et al., 2012), CC BY-SA 4.0. Model: google/siglip2-base-patch16-256, Apache-2.0.