in SAM2-Demo/Common/SAM2.swift [128:150]
func bestMask(for output: SAM2_1SmallMaskDecoderFLOAT16Output) -> MLMultiArray {
if #available(macOS 15.0, *) {
let scores = output.scoresShapedArray.scalars
let argmax = scores.firstIndex(of: scores.max() ?? 0) ?? 0
return MLMultiArray(output.low_res_masksShapedArray[0, argmax])
} else {
// Convert scores to float32 for compatibility with macOS < 15,
// plus ugly loop copy (could do some memcpys)
let scores = output.scores
let floatScores = (0..<scores.count).map { scores[$0].floatValue }
let argmax = floatScores.firstIndex(of: floatScores.max() ?? 0) ?? 0
let allMasks = output.low_res_masks
let (h, w) = (allMasks.shape[2], allMasks.shape[3])
let slice = try! MLMultiArray(shape: [h, w], dataType: allMasks.dataType)
for i in 0..<h.intValue {
for j in 0..<w.intValue {
let position = [0, argmax, i, j] as [NSNumber]
slice[[i as NSNumber, j as NSNumber]] = allMasks[position]
}
}
return slice
}
}