import { Tensor } from 'onnxruntime-web' ; import { ShapeUtil , TypedArrayUtil } from './yoloPostprocessUtils' ; export function concat ( x , axis ) { const input0 = x [ 0 ] ; const inputShape = input0 . dims ? input0 . dims : [ input0 . data . length ] ; if ( axis >= inputShape . length || axis < ( - 1 * inputShape . length ) ) { throw new Error ( `axis specified for concat doesn't match input dimensionality` ) ; } if ( axis < 0 ) { axis = inputShape . length + axis ; } let concatAxisSize = inputShape [ axis ] ; const outputShape = new Array < number > ( inputShape . length ) ; for ( let i = 1 ; i < x . length ; i ++ ) { const dataN = x [ i ] ; const dataNShape = dataN . dims ? dataN . dims : [ dataN . data . length ] ; for ( let axisIndex = 0 ; axisIndex < inputShape . length ; axisIndex ++ ) { if ( axisIndex === axis ) { concatAxisSize += dataNShape [ axisIndex ] ; } if ( inputShape [ axisIndex ] !== dataNShape [ axisIndex ] ) { throw new Error ( `non concat dimensions must match` ) ; } outputShape [ axisIndex ] = dataNShape [ axisIndex ] ; } } outputShape [ axis ] = concatAxisSize ; const output = new Tensor ( input0 . type , TypedArrayUtil . createTypedArray ( x [ 0 ] . type , ShapeUtil . size ( outputShape ) ) , outputShape ) ; const Y = output . data ; let axisPitch = 1 ; for ( let i = outputShape . length - 1 ; i >= axis ; i -- ) { axisPitch *= outputShape [ i ] ; } let outputBase = 0 ; for ( let inputIndex = 0 ; inputIndex < x . length ; inputIndex ++ ) { const dataN = x [ inputIndex ] ; const dataNDims = dataN . dims ? dataN . dims : [ dataN . data . length ] ; let inputAxisPitch = 1 ; for ( let i = dataNDims . length - 1 ; i >= axis ; i -- ) { inputAxisPitch *= dataNDims [ i ] ; } const inputData = dataN . data ; const inputSize = ShapeUtil . size ( dataNDims ) ; let outputOffset = outputBase ; for ( let i = 0 , j = 0 ; i < inputSize ; i ++ ) { Y [ outputOffset + i ] = inputData [ i ] ; if ( ++ j === inputAxisPitch ) { outputOffset += ( axisPitch - inputAxisPitch ) ; j = 0 ; } } outputBase += inputAxisPitch ; } return output ; }