license import { KernelConfig , KernelFunc , NumericDataType , TensorInfo , Transform , TransformAttrs , TransformInputs , TypedArray , util } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; export function transform ( args ) { const { inputs , attrs , backend } = args ; const { image , transforms } = inputs ; const { interpolation , fillMode , fillValue , outputShape } = attrs ; const [ batch , imageHeight , imageWidth , numChannels ] = image . shape ; const [ outHeight , outWidth ] = outputShape != null ? outputShape : [ imageHeight , imageWidth ] ; const outShape = [ batch , outHeight , outWidth , numChannels ] ; const strides = util . computeStrides ( image . shape ) ; const batchStride = strides [ 0 ] ; const rowStride = strides [ 1 ] ; const colStride = strides [ 2 ] ; const outVals = util . getTypedArrayFromDType ( image . dtype as NumericDataType , util . sizeFromShape ( outShape ) ) ; outVals . fill ( fillValue ) ; const imageVals = backend . data . get ( image . dataId ) . values as TypedArray ; const transformVals = backend . data . get ( transforms . dataId ) . values as TypedArray ; for ( let b = 0 ; b < batch ; ++ b ) { const transform = transforms . shape [ 0 ] === 1 ? transformVals : transformVals . subarray ( b * 8 , b * 8 + 8 ) ; for ( let outY = 0 ; outY < outHeight ; ++ outY ) { for ( let outX = 0 ; outX < outWidth ; ++ outX ) { for ( let channel = 0 ; channel < numChannels ; ++ channel ) { let val ; const projection = transform [ 6 ] * outX + transform [ 7 ] * outY + 1 ; if ( projection === 0 ) { continue ; } const inX = ( transform [ 0 ] * outX + transform [ 1 ] * outY + transform [ 2 ] ) / projection ; const inY = ( transform [ 3 ] * outX + transform [ 4 ] * outY + transform [ 5 ] ) / projection ; const x = mapCoord ( inX , imageWidth , fillMode ) ; const y = mapCoord ( inY , imageHeight , fillMode ) ; switch ( interpolation ) { case 'nearest' : val = nearestInterpolation ( imageVals , imageHeight , imageWidth , batchStride , rowStride , colStride , b , y , x , channel , fillValue ) ; break ; case 'bilinear' : val = bilinearInterpolation ( imageVals , imageHeight , imageWidth , batchStride , rowStride , colStride , b , y , x , channel , fillValue ) ; break ; default : throw new Error ( `Error in Transform: Expect 'nearest' or ` + `'bilinear', but got ${ interpolation }` ) ; } const ind = b * batchStride + outY * rowStride + outX * colStride + channel ; outVals [ ind ] = val ; } } } return backend . makeTensorInfo ( outShape , image . dtype , outVals ) ; } const dataId = backend . write ( outVals , outShape , image . dtype ) ; return { dataId , shape : image . shape , dtype : image . dtype } ; } export const transformConfig = { kernelName : Transform , backendName : 'cpu' , kernelFunc : transform as { } as KernelFunc } ; function mapCoord ( outCoord , len , mode ) { switch ( mode ) { case 'reflect' : return mapCoordReflect ( outCoord , len ) ; case 'wrap' : return mapCoordWrap ( outCoord , len ) ; case 'nearest' : return mapCoordNearest ( outCoord , len ) ; case 'constant' : default : return mapCoordConstant ( outCoord , len ) ; } } function mapCoordReflect ( outCoord , len ) { let inCoord = outCoord ; if ( inCoord < 0 ) { if ( len <= 1 ) { inCoord = 0 ; } else { const sz2 = 2 * len ; if ( inCoord < sz2 ) { inCoord = sz2 * Math . trunc ( - inCoord / sz2 ) + inCoord ; } inCoord = inCoord < - len ? inCoord + sz2 : - inCoord - 1 ; } } else if ( inCoord > len - 1 ) { if ( len <= 1 ) { inCoord = 0 ; } else { const sz2 = 2 * len ; inCoord -= sz2 * Math . trunc ( inCoord / sz2 ) ; if ( inCoord >= len ) { inCoord = sz2 - inCoord - 1 ; } } } return util . clamp ( 0 , inCoord , len - 1 ) ; } function mapCoordWrap ( outCoord , len ) { let inCoord = outCoord ; if ( inCoord < 0 ) { if ( len <= 1 ) { inCoord = 0 ; } else { const sz = len - 1 ; inCoord += len * ( Math . trunc ( - inCoord / sz ) + 1 ) ; } } else if ( inCoord > len - 1 ) { if ( len <= 1 ) { inCoord = 0 ; } else { const sz = len - 1 ; inCoord -= len * Math . trunc ( inCoord / sz ) ; } } return util . clamp ( 0 , inCoord , len - 1 ) ; } function mapCoordConstant ( outCoord , len ) { return outCoord ; } function mapCoordNearest ( outCoord , len ) { return util . clamp ( 0 , outCoord , len - 1 ) ; } function readWithFillValue ( imageVals , imageHeight , imageWidth , batchStride , rowStride , colStride , batch , y , x , channel , fillValue ) { const ind = batch * batchStride + y * rowStride + x * colStride + channel ; if ( 0 <= y && y < imageHeight && 0 <= x && x < imageWidth ) { return imageVals [ ind ] ; } else { return fillValue ; } } function nearestInterpolation ( imageVals , imageHeight , imageWidth , batchStride , rowStride , colStride , batch , y , x , channel , fillValue ) { const $y = Math . round ( y ) ; const $x = Math . round ( x ) ; return readWithFillValue ( imageVals , imageHeight , imageWidth , batchStride , rowStride , colStride , batch , $y , $x , channel , fillValue ) ; } function bilinearInterpolation ( imageVals , imageHeight , imageWidth , batchStride , rowStride , colStride , batch , y , x , channel , fillValue ) { const yFloor = Math . floor ( y ) ; const xFloor = Math . floor ( x ) ; const yCeil = yFloor + 1 ; const xCeil = xFloor + 1 ; const valueYFloor = ( xCeil - x ) * readWithFillValue ( imageVals , imageHeight , imageWidth , batchStride , rowStride , colStride , batch , yFloor , xFloor , channel , fillValue ) + ( x - xFloor ) * readWithFillValue ( imageVals , imageHeight , imageWidth , batchStride , rowStride , colStride , batch , yFloor , xCeil , channel , fillValue ) ; const valueYCeil = ( xCeil - x ) * readWithFillValue ( imageVals , imageHeight , imageWidth , batchStride , rowStride , colStride , batch , yCeil , xFloor , channel , fillValue ) + ( x - xFloor ) * readWithFillValue ( imageVals , imageHeight , imageWidth , batchStride , rowStride , colStride , batch , yCeil , xCeil , channel , fillValue ) ; return ( yCeil - y ) * valueYFloor + ( y - yFloor ) * valueYCeil ; }