diff --git a/Package.swift b/Package.swift index 60508acd..1e80f9b0 100644 --- a/Package.swift +++ b/Package.swift @@ -35,6 +35,12 @@ let package = Package( "CoreAIImageSegmenter" ] ), + .library( + name: "CoreAIVideoSegmentation", + targets: [ + "CoreAIVideoSegmenter" + ] + ), .library( name: "CoreAISpeech", targets: ["CoreAISpeech"] @@ -86,6 +92,14 @@ let package = Package( .enableUpcomingFeature("MemberImportVisibility") ] ), + .target( + name: "CoreAIVideoSegmenter", + dependencies: ["CoreAIShared"], + path: "swift/Sources/CoreAIVideoSegmenter", + swiftSettings: [ + .enableUpcomingFeature("MemberImportVisibility") + ] + ), // Shared utilities .target( @@ -208,6 +222,18 @@ let package = Package( .enableUpcomingFeature("MemberImportVisibility") ] ), + .executableTarget( + name: "video-segmenter", + dependencies: [ + "CoreAIVideoSegmenter", + "CoreAIShared", + .product(name: "ArgumentParser", package: "swift-argument-parser"), + ], + path: "swift/Sources/Tools/video-segmenter", + swiftSettings: [ + .enableUpcomingFeature("MemberImportVisibility") + ] + ), .executableTarget( name: "diffusion-runner", dependencies: [ @@ -297,6 +323,14 @@ let package = Package( ], path: "swift/Tests/ImageSegmenterTests" ), + .testTarget( + name: "VideoSegmenterTests", + dependencies: [ + "CoreAIVideoSegmenter", + "CoreAIShared", + ], + path: "swift/Tests/VideoSegmenterTests" + ), .testTarget( name: "DiffusionPipelineTests", dependencies: [ diff --git a/models/sam3_video/README.md b/models/sam3_video/README.md index a97038c1..134ab0a9 100644 --- a/models/sam3_video/README.md +++ b/models/sam3_video/README.md @@ -80,7 +80,7 @@ the whole video. | `tracker_encode` | no | | `text_encode` | once per prompt, per video | -## The Swift runtime (PENDING: this is design only) +## The Swift runtime `swift/Sources/CoreAIVideoSegmenter` mimics the HF logic in Swift, and the `video-segmenter` tool drives it end to end: diff --git a/python/src/coreai_models/segmentation/video_pipeline.py b/python/src/coreai_models/segmentation/video_pipeline.py index e96f471e..689065c3 100644 --- a/python/src/coreai_models/segmentation/video_pipeline.py +++ b/python/src/coreai_models/segmentation/video_pipeline.py @@ -715,7 +715,7 @@ async def _async_export_video(config: VideoExportConfig) -> str: logger.info("Saved Core AI asset to %s", asset_path) # Metadata before tokenizer, so a flaky HF fetch can't leave an unloadable bundle. - _write_bundle_metadata(bundle_dir, asset_path.name, config) + _write_bundle_metadata(bundle_dir, asset_path.name, config, model.config) _write_tokenizer(bundle_dir / "tokenizer", config.hf_model_id) return str(bundle_dir) @@ -841,14 +841,65 @@ def _resolve_paths(config: VideoExportConfig) -> tuple[Path, Path]: return bundle_dir, bundle_dir / f"{name}.aimodel" +#: ``Sam3VideoConfig`` fields the host runtime needs, copied into the bundle's ``tracking`` +#: block. None affect the traced graphs; they all govern host-side heuristics. A runtime +#: that hardcodes the upstream defaults only diverges on a checkpoint that tuned them. +_TRACKING_FIELDS = ( + "score_threshold_detection", + "det_nms_thresh", + "new_det_thresh", + "assoc_iou_thresh", + "trk_assoc_iou_thresh", + "high_conf_thresh", + "high_iou_thresh", + "recondition_every_nth_frame", + "recondition_on_trk_masks", + "hotstart_delay", + "hotstart_unmatch_thresh", + "hotstart_dup_thresh", + "suppress_unmatched_only_within_hotstart", + "init_trk_keep_alive", + "max_trk_keep_alive", + "min_trk_keep_alive", + "decrease_trk_keep_alive_for_empty_masklets", + "suppress_overlapping_based_on_recent_occlusion_threshold", + "max_num_objects", + "fill_hole_area", +) + +#: Tracker-config fields that decide which stored frames are eligible for the memory bank. +#: The export pins their sum (``spatial_slots`` is ``max_cond_frame_num + num_maskmem - 1``) +#: but not the split, so the host cannot recover them from the asset alone. +_TRACKER_MEMORY_FIELDS = ( + "num_maskmem", + "max_cond_frame_num", + "max_object_pointers_in_encoder", +) + + +def _tracking_metadata(config) -> dict: + """Collect the host-side thresholds from a ``Sam3VideoConfig``.""" + tracking: dict = {} + for field in _TRACKING_FIELDS: + if hasattr(config, field): + tracking[field] = getattr(config, field) + tracker_config = config.tracker_config + for field in _TRACKER_MEMORY_FIELDS: + if hasattr(tracker_config, field): + tracking[field] = getattr(tracker_config, field) + return tracking + + def _write_bundle_metadata( - bundle_dir: Path, asset_filename: str, config: VideoExportConfig + bundle_dir: Path, asset_filename: str, config: VideoExportConfig, model_config ) -> None: """Write the bundle manifest. ``runtime`` carries the slot geometry because the host has to pack memory to exactly the shapes the graph was traced with; deriving it from the HF config at load time would silently break if the export used non-default slots. + + ``tracking`` carries the checkpoint's own heuristic thresholds; see ``_TRACKING_FIELDS``. """ metadata = { "metadata_version": "0.2", @@ -861,6 +912,7 @@ def _write_bundle_metadata( "ptr_slots": config.ptr_slots, "max_text_seq_len": config.max_text_seq_len, }, + "tracking": _tracking_metadata(model_config), } metadata_path = bundle_dir / "metadata.json" with open(metadata_path, "w") as fh: diff --git a/python/tests/test_model_units/test_export/test_video_export.py b/python/tests/test_model_units/test_export/test_video_export.py index 6de9bbc5..b04f9884 100644 --- a/python/tests/test_model_units/test_export/test_video_export.py +++ b/python/tests/test_model_units/test_export/test_video_export.py @@ -7,11 +7,7 @@ Everything here runs on a randomly-initialized, heavily downscaled config (112x112, 2 backbone layers, 1 memory-attention layer) so no weights are -downloaded and the whole file runs in seconds. Numerical parity against the -real checkpoint is the parity harness's job -(``models/sam3_video/run_video_parity.py``); what these tests pin is the -*structure*: that the fixed-slot memory bank is mathematically equivalent to -HF's variable-length one, and that every entrypoint is traceable. +downloaded and the whole file runs in seconds. """ from __future__ import annotations diff --git a/swift/Sources/CoreAIShared/Bundle/BundleKind.swift b/swift/Sources/CoreAIShared/Bundle/BundleKind.swift index c5325abc..d2260403 100644 --- a/swift/Sources/CoreAIShared/Bundle/BundleKind.swift +++ b/swift/Sources/CoreAIShared/Bundle/BundleKind.swift @@ -13,5 +13,8 @@ public enum BundleKind: String, Codable, Sendable, CaseIterable { case vlm case diffusion case segmenter + /// Text-promptable video segmentation (SAM 3 video). Separate from `segmenter` + /// because the bundle carries a `runtime` block with the memory-bank geometry. + case videoSegmenter = "video_segmenter" case speechRecognizer = "speech_recognizer" } diff --git a/swift/Sources/CoreAIShared/Image/OverlayPalette.swift b/swift/Sources/CoreAIShared/Image/OverlayPalette.swift index 831b0013..553f8228 100644 --- a/swift/Sources/CoreAIShared/Image/OverlayPalette.swift +++ b/swift/Sources/CoreAIShared/Image/OverlayPalette.swift @@ -33,6 +33,16 @@ public enum OverlayPalette { return hsvToRGB(h: Float(index) / Float(count), s: 0.85, v: 0.95) } + /// Stable color for a tracked object id. + /// + /// Steps the hue by the golden ratio's fractional part rather than dividing the wheel + /// by a count: the total is unknown mid-video, and consecutive ids, which is what new + /// tracks get, stay far apart instead of nearly on top of each other. + public static func color(forID id: Int) -> (UInt8, UInt8, UInt8) { + let hue = (Float(id) * 0.618_033_99).truncatingRemainder(dividingBy: 1.0) + return hsvToRGB(h: hue < 0 ? hue + 1 : hue, s: 0.85, v: 0.95) + } + /// HSV → RGB, all components in [0, 1]. public static func hsvToRGB(h: Float, s: Float, v: Float) -> (UInt8, UInt8, UInt8) { let h6 = h * 6 diff --git a/swift/Sources/CoreAIShared/Runtime/BilinearResampler.swift b/swift/Sources/CoreAIShared/Runtime/BilinearResampler.swift new file mode 100644 index 00000000..498fd58e --- /dev/null +++ b/swift/Sources/CoreAIShared/Runtime/BilinearResampler.swift @@ -0,0 +1,370 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import Accelerate +import Foundation + +/// Separable bilinear resampling of a planar `Float` image, matching +/// `torch.nn.functional.interpolate(mode: "bilinear", align_corners: false)`. +/// +/// Construct once per (source, destination) pair and reuse; the weight tables are the only +/// setup cost. Per-frame callers should use `resample(_:into:scratch:)` — the allocating +/// overload is a convenience, not the fast path. +/// +/// `antialias` widens the filter when downsampling, the way PIL and +/// `torchvision.transforms.Resize` do. Upsampling ignores it. +public struct BilinearResampler: Sendable { + public let sourceWidth: Int + public let sourceHeight: Int + public let destinationWidth: Int + public let destinationHeight: Int + + private let horizontal: Weights + private let vertical: Weights + + public init( + sourceWidth: Int, sourceHeight: Int, + destinationWidth: Int, destinationHeight: Int, + antialias: Bool = false + ) { + precondition( + sourceWidth > 0 && sourceHeight > 0 && destinationWidth > 0 && destinationHeight > 0, + "BilinearResampler dimensions must be positive") + self.sourceWidth = sourceWidth + self.sourceHeight = sourceHeight + self.destinationWidth = destinationWidth + self.destinationHeight = destinationHeight + // Only the horizontal axis needs a control vector; the vertical pass walks whole + // contiguous rows and reads `starts`/`values` directly. + self.horizontal = Weights( + sourceSize: sourceWidth, destinationSize: destinationWidth, antialias: antialias, + buildControlVector: true) + self.vertical = Weights( + sourceSize: sourceHeight, destinationSize: destinationHeight, antialias: antialias, + buildControlVector: false) + } + + /// True when the resampler would copy its input unchanged. + public var isIdentity: Bool { + sourceWidth == destinationWidth && sourceHeight == destinationHeight + } + + /// True when the horizontal pass has to go the long way round. See `resampleTransposing`. + /// + /// This is the common case. Only a two-tap resize whose ratio is a binary fraction avoids + /// it, which in practice means the tracker's 4x mask upsamples on a patch-16 export; a + /// patch-14 export scales by 3.5 and transposes throughout. + private var transposes: Bool { horizontal.maxTaps > 2 || !horizontal.controlIsExact } + + /// Elements of caller-owned storage `resample(_:into:scratch:)` needs. + public var scratchCount: Int { + if isIdentity { return 0 } + if transposes { + // Vertical output, its transpose, and the horizontal output before transposing back. + return 2 * destinationHeight * sourceWidth + destinationWidth * destinationHeight + } + return sourceHeight * destinationWidth + } + + /// Resample a row-major `sourceHeight × sourceWidth` buffer into caller-owned storage. + /// + /// `scratch` must hold at least `scratchCount` elements and may be reused across calls; + /// its contents are not meaningful afterwards. This is the overload the per-frame paths + /// use, since at video resolution every buffer here is megabytes. + public func resample( + _ source: UnsafeBufferPointer, + into destination: UnsafeMutableBufferPointer, + scratch: UnsafeMutableBufferPointer + ) { + precondition( + source.count >= sourceWidth * sourceHeight, + "BilinearResampler expected \(sourceWidth * sourceHeight) samples, got \(source.count)") + precondition( + destination.count >= destinationWidth * destinationHeight, + "BilinearResampler needs a \(destinationWidth * destinationHeight) element " + + "destination, got \(destination.count)") + let input = source.baseAddress! + let output = destination.baseAddress! + if isIdentity { + output.update(from: input, count: sourceWidth * sourceHeight) + return + } + precondition( + scratch.count >= scratchCount, + "BilinearResampler needs \(scratchCount) elements of scratch, got \(scratch.count)") + let workspace = scratch.baseAddress! + + if transposes { + resampleTransposing(input: input, output: output, workspace: workspace) + return + } + + // Horizontal first, so the vertical pass runs over the smaller of the two widths + // when downscaling and the taps stay whole contiguous rows either way. + horizontal.applyAcrossRows( + input: input, inputStride: sourceWidth, + output: workspace, outputStride: destinationWidth, + rowCount: sourceHeight) + vertical.applyDownColumns( + input: workspace, output: output, rowWidth: destinationWidth) + } + + /// The path for everything `vDSP_vlint` cannot do exactly: an antialiased downsample, + /// whose widened filter reads a contiguous run of source columns from an irregularly + /// spaced start, and any resize whose control vector would lose precision, per + /// `Weights.controlIsExact`. + /// + /// Transposing turns the horizontal taps back into whole contiguous rows, so both axes + /// run through `applyDownColumns`. Vertical goes first because whichever axis shrinks + /// should shrink before the transposes have to move it. + private func resampleTransposing( + input: UnsafePointer, + output: UnsafeMutablePointer, + workspace: UnsafeMutablePointer + ) { + let rows = workspace + let transposed = rows + destinationHeight * sourceWidth + let columns = transposed + sourceWidth * destinationHeight + + vertical.applyDownColumns(input: input, output: rows, rowWidth: sourceWidth) + // `vDSP_mtrans(A, _, C, _, M, N)` writes an M×N result from an N×M input. + vDSP_mtrans( + rows, 1, transposed, 1, vDSP_Length(sourceWidth), vDSP_Length(destinationHeight)) + horizontal.applyDownColumns( + input: transposed, output: columns, rowWidth: destinationHeight) + vDSP_mtrans( + columns, 1, output, 1, vDSP_Length(destinationHeight), + vDSP_Length(destinationWidth)) + } + + /// Resample into a caller-owned destination, reusing `scratch` across calls. + /// + /// The array-level form of `resample(_:into:scratch:)`. Per-frame callers keep both + /// buffers alive and pay no allocation here. + public func resample( + _ source: [Float], into destination: inout [Float], scratch: inout [Float] + ) { + source.withUnsafeBufferPointer { input in + destination.withUnsafeMutableBufferPointer { output in + scratch.withUnsafeMutableBufferPointer { scratch in + resample(input, into: output, scratch: scratch) + } + } + } + } + + /// Resample a row-major `sourceHeight × sourceWidth` buffer. + /// + /// Allocates both the result and the scratch on every call. Fine for one-shot use; use + /// `resample(_:into:scratch:)` on anything that runs per frame. + public func resample(_ source: [Float]) -> [Float] { + var destination = [Float](repeating: 0, count: destinationHeight * destinationWidth) + var scratch = [Float](repeating: 0, count: scratchCount) + resample(source, into: &destination, scratch: &scratch) + return destination + } +} + +// MARK: - Weights + +/// Per-output-index filter taps along one axis. +/// +/// Stored as a dense `destinationSize × maxTaps` table with a start index and live tap +/// count per output. Rows near an edge use fewer taps than the interior; padding to a +/// fixed stride keeps the indexing uniform at the cost of a few unused slots. +private struct Weights: Sendable { + let destinationSize: Int + let maxTaps: Int + /// First source index contributing to each output index. + let starts: [Int] + /// Number of live taps for each output index. + let counts: [Int] + /// `destinationSize * maxTaps` coefficients, normalized to sum to 1 per output. + let values: [Float] + /// Source coordinate per output, the control vector `vDSP_vlint` interpolates from: the + /// same taps as `values`, in the form Accelerate can apply a whole row at a time. Built + /// only for the horizontal axis of a two-tap resize. + let positions: [Float] + /// First output whose second tap falls off the end of the row. See `applyAcrossRows`. + let clampedFrom: Int + /// True when `positions` drives `vDSP_vlint` to the same taps the `Double` weights carry. + /// + /// `vDSP_vlint` recovers the interpolation fraction from a single `Float` coordinate, so + /// its precision is set by the magnitude of that coordinate rather than by the fraction. + /// At video widths the fraction quantizes coarsely enough to shift a 0-255 sample by a + /// few hundredths of a code value — enough to flip `FramePreprocessor`'s rounding and + /// drift from torch. The error is exactly zero when the resize ratio is a binary + /// fraction, so this is checked per resize rather than assumed either way. + let controlIsExact: Bool + + init(sourceSize: Int, destinationSize: Int, antialias: Bool, buildControlVector: Bool) { + self.destinationSize = destinationSize + let scale = Double(sourceSize) / Double(destinationSize) + + if antialias && scale > 1.0 { + // Downsampling with a triangle filter whose support grows with the ratio. + // Mirrors `aten/src/ATen/native/UpSample.h::_compute_weights_aa`. + let support = scale + let inverseScale = 1.0 / scale + let taps = Int(support.rounded(.up)) * 2 + 1 + self.maxTaps = taps + var starts = [Int](repeating: 0, count: destinationSize) + var counts = [Int](repeating: 0, count: destinationSize) + var values = [Float](repeating: 0, count: destinationSize * taps) + for index in 0..= Float(sourceSize - 1) } ?? destinationSize) + : destinationSize + } + + /// Resample each row independently, two Accelerate calls per row. + /// + /// Two-tap only: `vDSP_vlint` interpolates straight from a source coordinate, which is + /// both the weight table and the inner loop. Multi-tap resizes transpose instead and use + /// `applyDownColumns`. + func applyAcrossRows( + input: UnsafePointer, inputStride: Int, + output: UnsafeMutablePointer, outputStride: Int, + rowCount: Int + ) { + precondition(maxTaps == 2, "applyAcrossRows is the two-tap path") + let interpolated = vDSP_Length(clampedFrom) + let clamped = vDSP_Length(destinationSize - clampedFrom) + let sourceLength = vDSP_Length(inputStride) + positions.withUnsafeBufferPointer { positions in + for row in 0.. 0 { + vDSP_vlint( + sourceRow, positions.baseAddress!, 1, destinationRow, 1, + interpolated, sourceLength) + } + if clamped > 0 { + vDSP_vfill( + sourceRow + inputStride - 1, destinationRow + clampedFrom, 1, clamped) + } + } + } + } + + /// Resample down the columns. Each tap is a whole contiguous row scaled by one + /// coefficient, so this is a handful of vDSP calls per output row rather than a + /// per-pixel loop. + func applyDownColumns( + input: UnsafePointer, + output: UnsafeMutablePointer, + rowWidth: Int + ) { + let length = vDSP_Length(rowWidth) + for index in 0.. 1 else { + // Clamped edge: the table folds both weights onto one sample, so the + // coefficient is exactly 1 and this is a copy. + destinationRow.update(from: sourceRow, count: rowWidth) + continue + } + // `A + f * (B - A)` in one pass, where vsmul plus vsma is two passes and two + // writes over what is the memory-bound axis when upsampling. + var fraction = values[base + 1] + vDSP_vintb( + sourceRow, 1, sourceRow + rowWidth, 1, &fraction, destinationRow, 1, length) + continue + } + + var first = values[base] + vDSP_vsmul(sourceRow, 1, &first, destinationRow, 1, length) + guard counts[index] > 1 else { continue } + for tap in 1..= 0 && height >= 0, "MaskBitset dimensions must be non-negative") + self.width = width + self.height = height + self.words = [UInt64](repeating: 0, count: (width * height + 63) / 64) + } + + /// Threshold a row-major float buffer: `value > threshold` becomes a set bit. + /// + /// The comparison is strictly greater to match HF, which binarizes mask logits with + /// `mask > 0` throughout; `>=` would flip every exactly-zero pixel. + public init(thresholding values: [Float], width: Int, height: Int, above threshold: Float = 0) { + self.init(width: width, height: height) + precondition( + values.count >= width * height, + "MaskBitset needs \(width * height) values, got \(values.count)") + let count = width * height + words.withUnsafeMutableBufferPointer { output in + values.withUnsafeBufferPointer { input in + var index = 0 + var wordIndex = 0 + while index < count { + let end = min(index + 64, count) + var word: UInt64 = 0 + var bit: UInt64 = 1 + for i in index.. threshold { word |= bit } + bit <<= 1 + } + output[wordIndex] = word + wordIndex += 1 + index = end + } + } + } + } + + /// Threshold a sub-range of a larger row-major buffer, one mask out of a stacked + /// `[N, H, W]` tensor. + public init( + thresholding values: ArraySlice, width: Int, height: Int, above threshold: Float = 0 + ) { + self.init(thresholding: Array(values), width: width, height: height, above: threshold) + } + + public subscript(x: Int, y: Int) -> Bool { + get { + let index = y * width + x + return words[index >> 6] & (1 << UInt64(index & 63)) != 0 + } + set { + let index = y * width + x + if newValue { + words[index >> 6] |= 1 << UInt64(index & 63) + } else { + words[index >> 6] &= ~(1 << UInt64(index & 63)) + } + } + } + + /// Number of set pixels. + public var area: Int { + var total = 0 + for word in words { total += word.nonzeroBitCount } + return total + } + + public var isEmpty: Bool { + for word in words where word != 0 { return false } + return true + } + + /// Intersection over union, matching `modeling_sam3_video.mask_iou`. + /// + /// HF clamps the union to a minimum of 1, so empty-vs-empty scores 0 rather than + /// dividing by zero. Association therefore treats two empty masks as unrelated. + public func iou(_ other: MaskBitset) -> Float { + precondition( + width == other.width && height == other.height, + "MaskBitset.iou requires matching dimensions") + var intersection = 0 + var union = 0 + for index in words.indices { + let a = words[index] + let b = other.words[index] + intersection += (a & b).nonzeroBitCount + union += (a | b).nonzeroBitCount + } + return Float(intersection) / Float(max(union, 1)) + } + + /// Tight bounding box of the set pixels, top-left origin. Empty masks give `.zero`. + /// + /// `torchvision.ops.masks_to_boxes` reports inclusive extremes, so a single set pixel is + /// a zero-sized rect. Kept that way so the values compare directly against a reference. + public var boundingBox: CGRect { + var minX = width + var minY = height + var maxX = -1 + var maxY = -1 + forEachSetIndex { index in + let y = index / width + let x = index - y * width + if x < minX { minX = x } + if x > maxX { maxX = x } + if y < minY { minY = y } + if y > maxY { maxY = y } + } + guard maxX >= 0 else { return .zero } + return CGRect( + x: CGFloat(minX), y: CGFloat(minY), + width: CGFloat(maxX - minX), height: CGFloat(maxY - minY)) + } + + /// Visit every set pixel's row-major index in ascending order. + @inline(__always) + public func forEachSetIndex(_ body: (Int) -> Void) { + let count = width * height + for wordIndex in words.indices { + var word = words[wordIndex] + let base = wordIndex << 6 + while word != 0 { + let index = base + Int(word.trailingZeroBitCount) + if index >= count { return } + body(index) + word &= word - 1 + } + } + } + + /// Row-major bytes, one per pixel, 1 for foreground. + public func toBytes() -> [UInt8] { + var out = [UInt8](repeating: 0, count: width * height) + forEachSetIndex { out[$0] = 1 } + return out + } + + /// Rebuild from `numpy.packbits` output: MSB-first within each byte, row-major. + public init(packedBits: [UInt8], width: Int, height: Int) { + self.init(width: width, height: height) + let count = width * height + for index in 0..> 3] + if byte & (0x80 >> UInt8(index & 7)) != 0 { + words[index >> 6] |= 1 << UInt64(index & 63) + } + } + } +} diff --git a/swift/Sources/CoreAIShared/Runtime/ModelStructure.swift b/swift/Sources/CoreAIShared/Runtime/ModelStructure.swift index 834a499e..68f2f39c 100644 --- a/swift/Sources/CoreAIShared/Runtime/ModelStructure.swift +++ b/swift/Sources/CoreAIShared/Runtime/ModelStructure.swift @@ -17,6 +17,12 @@ public enum GraphNames { public static let imageEncode = "image_encode" public static let textEncode = "text_encode" public static let detect = "detect" + // Video segmenter (SAM3 video export). Shares the three names above with the + // image segmenter, so `trackerStep` is what tells the two apart. + public static let trackerEncode = "tracker_encode" + public static let trackerStep = "tracker_step" + public static let memoryEncode = "memory_encode" + public static let trackerMaskInit = "tracker_mask_init" } /// Represents the detected structure of a Core AI model. @@ -38,6 +44,10 @@ public enum ModelStructure: Equatable, Sendable, CustomStringConvertible { /// Identified by presence of `image_encode`, `text_encode`, and `detect` graphs. case multiFunctionSegmenter + /// Seven-function SAM3 video segmenter. + /// Identified by `tracker_step`, which no other export produces. + case videoSegmenter + public var description: String { switch self { case .chunkedStatic(let batchSize): @@ -46,6 +56,8 @@ public enum ModelStructure: Equatable, Sendable, CustomStringConvertible { return "dynamic" case .multiFunctionSegmenter: return "multiFunctionSegmenter" + case .videoSegmenter: + return "videoSegmenter" } } @@ -54,11 +66,12 @@ public enum ModelStructure: Equatable, Sendable, CustomStringConvertible { /// - `chunkedStatic` → NeuralEngine /// - `dynamic` → GPU /// - `multiFunctionSegmenter` → NeuralEngine + /// - `videoSegmenter` → GPU public var preferredDevice: String { switch self { case .chunkedStatic, .multiFunctionSegmenter: return "NeuralEngine" - case .dynamic: + case .dynamic, .videoSegmenter: return "GPU" } } @@ -68,6 +81,7 @@ public enum ModelStructure: Equatable, Sendable, CustomStringConvertible { /// - `chunkedStatic` → prefer `.neuralEngine` /// - `dynamic` → prefer `.gpu` + `expectFrequentReshapes` /// - `multiFunctionSegmenter` → prefer `.neuralEngine` + /// - `videoSegmenter` → prefer `.gpu`, without `expectFrequentReshapes` public var specializationOptions: SpecializationOptions { switch self { case .chunkedStatic, .multiFunctionSegmenter: @@ -76,6 +90,9 @@ public enum ModelStructure: Equatable, Sendable, CustomStringConvertible { var opts = SpecializationOptions(preferredComputeUnitKind: .gpu) opts.expectFrequentReshapes = true return opts + case .videoSegmenter: + // TODO: the optimized export will need different specialization options. + return SpecializationOptions(preferredComputeUnitKind: .gpu) } } } @@ -247,6 +264,13 @@ public struct PreparedModel: Sendable { return .chunkedStatic(batchSize: batchSize) } + // Must be checked before the image segmenter below: this asset also declares + // image_encode / text_encode / detect, and the looser test would specialize it + // for the Neural Engine, which cannot compile two of its graphs. + if graphSet.contains(GraphNames.trackerStep) { + return .videoSegmenter + } + // Multi-function segmenter (e.g. optimized SAM3 — image_encode / text_encode / detect). // Targets neuralEngine; checked before the `main` fallback because some asset variants ship // a thin `main` graph alongside the trio. diff --git a/swift/Sources/CoreAIShared/Runtime/NDArray+Helpers.swift b/swift/Sources/CoreAIShared/Runtime/NDArray+Helpers.swift index 01c56348..b9ba19e1 100644 --- a/swift/Sources/CoreAIShared/Runtime/NDArray+Helpers.swift +++ b/swift/Sources/CoreAIShared/Runtime/NDArray+Helpers.swift @@ -160,6 +160,103 @@ public func fillFloatNDArray(_ array: inout NDArray, with elements: ArraySlice( + _ array: inout NDArray, as type: T.Type, elementOffset: Int, from source: [T] +) { + copyIntoNDArray(array.mutableRawView(), as: T.self, elementOffset: elementOffset, from: source) +} + +/// Copy `source` into an NDArray's MutableRawView starting at logical element `elementOffset`. +/// Takes the view as `consuming`, so the caller gives up ownership. +public func copyIntoNDArray( + _ rawView: consuming NDArray.MutableRawView, as type: T.Type, elementOffset: Int, + from source: [T] +) { + guard !source.isEmpty else { return } + writeNDArrayRegion( + rawView, as: type, label: "copyIntoNDArray", + elementOffset: elementOffset, count: source.count, + bulk: { destination, count in + source.withUnsafeBufferPointer { destination.update(from: $0.baseAddress!, count: count) } + }, + element: { source[$0] }) +} + +/// Write `value` into every logical element of `array` in `elementRange`. +/// +/// The offset-aware sibling of `fillNDArray(_:as:with:)`, which can only write from element 0. +public func fillNDArray( + _ array: inout NDArray, as type: T.Type, elementRange: Range, with value: T +) { + fillNDArray(array.mutableRawView(), as: type, elementRange: elementRange, with: value) +} + +/// MutableRawView overload of the region fill. +public func fillNDArray( + _ rawView: consuming NDArray.MutableRawView, as type: T.Type, elementRange: Range, + with value: T +) { + guard !elementRange.isEmpty else { return } + writeNDArrayRegion( + rawView, as: type, label: "fillNDArray", + elementOffset: elementRange.lowerBound, count: elementRange.count, + bulk: { destination, count in destination.update(repeating: value, count: count) }, + element: { _ in value }) +} + +/// Write `count` logical elements into `rawView` starting at element `elementOffset`. +/// +/// `bulk` takes the contiguous fast path; `element` feeds the stride walk used when the array +/// is padded for hardware alignment. +@inline(__always) +private func writeNDArrayRegion( + _ rawView: consuming NDArray.MutableRawView, as type: T.Type, label: StaticString, + elementOffset: Int, count: Int, + bulk: (UnsafeMutablePointer, Int) -> Void, + element: (Int) -> T +) { + let view = rawView.view(as: type) + view.withUnsafeMutablePointer { ptr, shape, strides in + let capacity = shape.product + precondition( + elementOffset >= 0 && elementOffset + count <= capacity, + "\(label): [\(elementOffset), \(elementOffset + count)) exceeds capacity \(capacity)") + + if isContiguousRowMajor(shape: shape, strides: strides) { + bulk(ptr + elementOffset, count) + return + } + + let rank = shape.count + var indices = [Int](repeating: 0, count: rank) + var remainder = elementOffset + for d in (0..= 0 { + indices[dim] += 1 + offset += strides[dim] + if indices[dim] < shape[dim] { break } + indices[dim] = 0 + offset -= strides[dim] * shape[dim] + dim -= 1 + } + } + } +} + // MARK: - Flatten Helpers /// Flatten an NDArray output into `[Float]`, branching on its own scalar type. diff --git a/swift/Sources/CoreAIShared/Text/CLIPTokenizer.swift b/swift/Sources/CoreAIShared/Text/CLIPTokenizer.swift index b450c1bd..be714dcb 100644 --- a/swift/Sources/CoreAIShared/Text/CLIPTokenizer.swift +++ b/swift/Sources/CoreAIShared/Text/CLIPTokenizer.swift @@ -110,6 +110,35 @@ public struct CLIPTokenizer: Sendable { return ids } + /// Encode `text` and report which slots hold real tokens. + /// + /// The pad token is `<|endoftext|>`, the same id that ends a real sequence, so the ids + /// alone can't tell padding from content. A prompt longer than `contextLength` is + /// truncated with `eotTokenId` forced into the last slot, making the mask all ones. + public func encodeWithMask( + _ text: String, contextLength: Int = 77 + ) -> (ids: [Int32], attentionMask: [Int32]) { + let cleaned = whitespaceClean(text).lowercased() + let wordTokens = tokenize(cleaned) + + var ids: [Int32] = [Self.sotTokenId] + ids += wordTokens.compactMap { encoder[$0] } + ids.append(Self.eotTokenId) + + if ids.count > contextLength { + ids = Array(ids.prefix(contextLength)) + ids[contextLength - 1] = Self.eotTokenId + } + + let realCount = ids.count + while ids.count < contextLength { + ids.append(Self.eotTokenId) + } + + let attentionMask = (0.. [String] { diff --git a/swift/Sources/CoreAIVideoSegmenter/FramePreprocessor.swift b/swift/Sources/CoreAIVideoSegmenter/FramePreprocessor.swift new file mode 100644 index 00000000..d131ec97 --- /dev/null +++ b/swift/Sources/CoreAIVideoSegmenter/FramePreprocessor.swift @@ -0,0 +1,137 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import Accelerate +import CoreAIShared +import CoreGraphics +import Foundation + +/// Turns a decoded frame into the planar CHW tensor `image_encode` expects. +/// +/// Deliberately not `CoreAIShared.ImagePreprocessor`, which the image segmenter uses: that +/// resizes by drawing through a `CGContext` at `.high` interpolation, a Lanczos-family kernel +/// on 8-bit samples, while `Sam3VideoVideoProcessor` resizes with plain bilinear in float. For +/// a single image the difference is invisible; across a tracked video it compounds, because +/// every frame's input feeds the memory bank. +/// +/// So: convert to float at native resolution first, resample in float, normalize last. +struct FramePreprocessor { + let targetSize: Int + let mean: (Float, Float, Float) + let standardDeviation: (Float, Float, Float) + + /// Resamplers are keyed by source size and rebuilt only when it changes, which for a + /// video is once. + private final class Cache { + var width = 0 + var height = 0 + var resampler: BilinearResampler? + var scratch: [Float] = [] + } + private let cache = Cache() + + init( + targetSize: Int, + mean: (CGFloat, CGFloat, CGFloat), + standardDeviation: (CGFloat, CGFloat, CGFloat) + ) { + self.targetSize = targetSize + self.mean = (Float(mean.0), Float(mean.1), Float(mean.2)) + self.standardDeviation = ( + Float(standardDeviation.0), Float(standardDeviation.1), Float(standardDeviation.2) + ) + } + + /// Preprocess a decoded frame. Returns flat `[3, targetSize, targetSize]`. + func preprocess(_ image: CGImage) throws -> [Float] { + let width = image.width + let height = image.height + guard let colorSpace = CGColorSpace(name: CGColorSpace.sRGB), + let context = CGContext( + data: nil, width: width, height: height, + bitsPerComponent: 8, bytesPerRow: width * 4, + space: colorSpace, + bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue), + let base = context.data + else { + throw ImagePreprocessorError.renderFailed + } + // Drawn at native size: this is a format conversion, not a resize. + context.draw(image, in: CGRect(x: 0, y: 0, width: width, height: height)) + let pixels = base.bindMemory(to: UInt8.self, capacity: width * height * 4) + return preprocess( + interleavedRGB: UnsafeBufferPointer(start: pixels, count: width * height * 4), + width: width, height: height, channelStride: 4) + } + + /// Preprocess interleaved 8-bit samples directly. + /// + /// `channelStride` is 4 for RGBA and 3 for packed RGB. The second form is what the + /// preprocessing tests feed in, so they can hold the decoder constant. + func preprocess( + interleavedRGB bytes: [UInt8], width: Int, height: Int, channelStride: Int + ) -> [Float] { + bytes.withUnsafeBufferPointer { + preprocess(interleavedRGB: $0, width: width, height: height, channelStride: channelStride) + } + } + + private func preprocess( + interleavedRGB bytes: UnsafeBufferPointer, + width: Int, height: Int, channelStride: Int + ) -> [Float] { + let resampler = resampler(sourceWidth: width, sourceHeight: height) + let sourcePixels = width * height + let targetPixels = targetSize * targetSize + + var output = [Float](repeating: 0, count: 3 * targetPixels) + var plane = [Float](repeating: 0, count: sourcePixels) + var resized = [Float](repeating: 0, count: targetPixels) + let means = [mean.0, mean.1, mean.2] + let deviations = [standardDeviation.0, standardDeviation.1, standardDeviation.2] + var elementCount = Int32(targetPixels) + + for channel in 0..<3 { + // De-interleave, staying in 0-255 so the rounding below lands on the same + // grid torchvision uses. + vDSP_vfltu8( + bytes.baseAddress! + channel, channelStride, &plane, 1, vDSP_Length(sourcePixels)) + + resampler.resample(plane, into: &resized, scratch: &cache.scratch) + + // torchvision resizes a uint8 tensor as uint8: it interpolates and then rounds + // back to integers. Skipping this leaves a uniform ~0.25-code-value bias against + // the reference. + vvnintf(&resized, resized, &elementCount) + + // Fold rescale and normalize into one affine pass: (x / 255 - m) / s. + var slope = 1 / (255 * deviations[channel]) + var offset = -means[channel] / deviations[channel] + output.withUnsafeMutableBufferPointer { out in + vDSP_vsmsa( + resized, 1, &slope, &offset, + out.baseAddress! + channel * targetPixels, 1, vDSP_Length(targetPixels)) + } + } + return output + } + + private func resampler(sourceWidth: Int, sourceHeight: Int) -> BilinearResampler { + if let existing = cache.resampler, cache.width == sourceWidth, cache.height == sourceHeight { + return existing + } + // `antialias: false` matches the video processor, and every real clip upscales to + // 1008 anyway, where the flag makes no difference. + let built = BilinearResampler( + sourceWidth: sourceWidth, sourceHeight: sourceHeight, + destinationWidth: targetSize, destinationHeight: targetSize, + antialias: false) + cache.width = sourceWidth + cache.height = sourceHeight + cache.resampler = built + cache.scratch = [Float](repeating: 0, count: built.scratchCount) + return built + } +} diff --git a/swift/Sources/CoreAIVideoSegmenter/Postprocessing/ConnectedComponents.swift b/swift/Sources/CoreAIVideoSegmenter/Postprocessing/ConnectedComponents.swift new file mode 100644 index 00000000..22b151a8 --- /dev/null +++ b/swift/Sources/CoreAIVideoSegmenter/Postprocessing/ConnectedComponents.swift @@ -0,0 +1,111 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import Foundation + +/// 8-connected component labelling and the mask cleanup built on it. +/// +/// Port of `fill_holes_in_mask_scores` and `_get_connected_components_with_padding`. Upstream +/// routes the labelling through an optional `kernels-community/cv-utils` CUDA kernel and, +/// when it is not installed, returns fake component areas of `H*W + 1`, which silently +/// no-ops both cleanup steps. This implementation is always present, so set +/// ``VideoSegmentationParameters/fillHoleArea`` to 0 to reproduce a Python run that was +/// missing the kernel. +/// +/// 8-connectivity matches SAM 2's `get_connected_components`. +enum ConnectedComponents { + /// Per-pixel component area for the set pixels of `mask`. Unset pixels get 0. + /// + /// Union-find with path halving over a single raster pass, then one pass to resolve roots + /// and a third to broadcast the root's area. Linear in pixels. + static func areas(of mask: [Bool], width: Int, height: Int) -> [Int32] { + let count = width * height + precondition(mask.count >= count, "ConnectedComponents.areas: mask is too small") + var parent = [Int32](repeating: -1, count: count) + + func find(_ start: Int32) -> Int32 { + var node = start + while parent[Int(node)] != node { + // Path halving: point each node at its grandparent while climbing. + parent[Int(node)] = parent[Int(parent[Int(node)])] + node = parent[Int(node)] + } + return node + } + func union(_ a: Int32, _ b: Int32) { + let rootA = find(a) + let rootB = find(b) + if rootA == rootB { return } + // Attach the larger index under the smaller so the labelling is deterministic. + if rootA < rootB { parent[Int(rootB)] = rootA } else { parent[Int(rootA)] = rootB } + } + + for y in 0.. 0, mask[index - 1] { union(Int32(index), Int32(index - 1)) } + if y > 0 { + let above = index - width + if mask[above] { union(Int32(index), Int32(above)) } + if x > 0, mask[above - 1] { union(Int32(index), Int32(above - 1)) } + if x + 1 < width, mask[above + 1] { union(Int32(index), Int32(above + 1)) } + } + } + } + + var componentArea = [Int32](repeating: 0, count: count) + var roots = [Int32](repeating: -1, count: count) + for index in 0..= 0 { + let root = find(Int32(index)) + roots[index] = root + componentArea[Int(root)] += 1 + } + var result = [Int32](repeating: 0, count: count) + for index in 0..= 0 { + result[index] = componentArea[Int(roots[index])] + } + return result + } + + /// Fill small background holes and remove small foreground specks, in place. + /// + /// Port of `fill_holes_in_mask_scores(mask, max_area, fill_holes=True, + /// remove_sprinkles=True)`. The two sentinel values (`0.1` and `-0.1`) are upstream's: + /// the mask stays a logit field, so a filled hole becomes weakly positive rather than + /// saturated. + static func fillHoles(_ logits: inout [Float], width: Int, height: Int, maxArea: Int) { + guard maxArea > 0 else { return } + let count = width * height + precondition(logits.count >= count, "fillHoles: logit buffer is too small") + + // Background: components of `logits <= 0` up to `maxArea` become weakly positive. + var background = [Bool](repeating: false, count: count) + for index in 0.. 0` up to the smaller of `maxArea` and half the + // mask's own area become weakly negative. Halving is what keeps a genuinely tiny + // object from deleting itself. + var foreground = [Bool](repeating: false, count: count) + var foregroundArea = 0 + for index in 0.. 0 { + foreground[index] = true + foregroundArea += 1 + } + let threshold = Int32(min(maxArea, foregroundArea / 2)) + guard threshold > 0 else { return } + let foregroundAreas = areas(of: foreground, width: width, height: height) + for index in 0..