license import { KernelConfig , NumericDataType , TypedArray } from '@tensorflow/tfjs-core' ; import { backend_util , RotateWithOffset , RotateWithOffsetAttrs , RotateWithOffsetInputs , util } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; export const rotateWithOffsetConfig = { kernelName : RotateWithOffset , backendName : 'cpu' , kernelFunc : ( { inputs , attrs , backend } ) => { const { image } = inputs as RotateWithOffsetInputs ; const { radians , fillValue , center } = attrs as { } as RotateWithOffsetAttrs ; const cpuBackend = backend as MathBackendCPU ; const output = util . getTypedArrayFromDType ( image . dtype as NumericDataType , util . sizeFromShape ( image . shape ) ) ; const [ batch , imageHeight , imageWidth , numChannels ] = image . shape ; const [ centerX , centerY ] = backend_util . getImageCenter ( center , imageHeight , imageWidth ) ; const fullOpacityValue = 255 ; const sinFactor = Math . sin ( radians ) ; const cosFactor = Math . cos ( radians ) ; const imageVals = cpuBackend . data . get ( image . dataId ) . values as TypedArray ; for ( let batchIdx = 0 ; batchIdx < batch ; batchIdx ++ ) { const batchOffset = batchIdx * imageWidth * imageHeight * numChannels ; for ( let row = 0 ; row < imageHeight ; row ++ ) { const rowOffset = row * ( imageWidth * numChannels ) ; for ( let col = 0 ; col < imageWidth ; col ++ ) { const colOffset = col * numChannels ; for ( let channel = 0 ; channel < numChannels ; channel ++ ) { const coords = [ batch , row , col , channel ] ; const x = coords [ 2 ] ; const y = coords [ 1 ] ; let coordX = ( x - centerX ) * cosFactor - ( y - centerY ) * sinFactor ; let coordY = ( x - centerX ) * sinFactor + ( y - centerY ) * cosFactor ; coordX = Math . round ( coordX + centerX ) ; coordY = Math . round ( coordY + centerY ) ; let outputValue = fillValue ; if ( typeof fillValue !== 'number' ) { if ( channel === 3 ) { outputValue = fullOpacityValue ; } else { outputValue = fillValue [ channel ] ; } } if ( coordX >= 0 && coordX < imageWidth && coordY >= 0 && coordY < imageHeight ) { const rotatedRowOffset = coordY * ( imageWidth * numChannels ) ; const rotatedColOffset = coordX * numChannels ; const imageIdx = batchOffset + rotatedRowOffset + rotatedColOffset + channel ; outputValue = imageVals [ imageIdx ] ; } const outIdx = batchOffset + rowOffset + colOffset + channel ; output [ outIdx ] = outputValue as number ; } } } } const dataId = cpuBackend . write ( output , image . shape , image . dtype ) ; return { dataId , shape : image . shape , dtype : image . dtype } ; } } ;