func bestMask()

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
        }
    }