license import { GPGPUProgram } from './gpgpu_math' ; export class CropAndResizeProgram implements GPGPUProgram { variableNames = [ 'Image' , 'Boxes' , 'BoxInd' ] ; outputShape : number [ ] = [ ] ; userCode : string ; constructor ( imageShape , boxShape , cropSize , method , extrapolationValue ) { const [ batch , imageHeight , imageWidth , depth ] = imageShape ; const [ numBoxes , ] = boxShape ; const [ cropHeight , cropWidth ] = cropSize ; this . outputShape = [ numBoxes , cropHeight , cropWidth , depth ] ; const methodId = method === 'bilinear' ? 1 : 0 ; const [ inputHeightFloat , inputWidthFloat ] = [ `${ imageHeight - 1 }.0` , `${ imageWidth - 1 }.0` ] ; const [ heightRatio , heightScale , inY ] = cropHeight > 1 ? [ `${ ( imageHeight - 1 ) / ( cropHeight - 1 ) }` , '(y2-y1) * height_ratio' , `y1*${ inputHeightFloat } + float(y)*(height_scale)` , ] : [ '0.0' , '0.0' , `0.5 * (y1+y2) * ${ inputHeightFloat }` , ] ; const [ widthRatio , widthScale , inX ] = cropWidth > 1 ? [ `${ ( imageWidth - 1 ) / ( cropWidth - 1 ) }` , '(x2-x1) * width_ratio' , `x1*${ inputWidthFloat } + float(x)*(width_scale)` , ] : [ '0.0' , '0.0' , `0.5 * (x1+x2) * ${ inputWidthFloat }` , ] ; this . userCode = `
      const float height_ratio = float(${ heightRatio });
      const float width_ratio = float(${ widthRatio });
      void main() {
        ivec4 coords = getOutputCoords();
        int b = coords[0];
        int y = coords[1];
        int x = coords[2];
        int d = coords[3];

        // get box vals
        float y1 = getBoxes(b,0);
        float x1 = getBoxes(b,1);
        float y2 = getBoxes(b,2);
        float x2 = getBoxes(b,3);

        // get image in batch index
        int bInd = round(getBoxInd(b));
        if(bInd < 0 || bInd >= ${ batch }) {
          return;
        }

        float height_scale = ${ heightScale };
        float width_scale = ${ widthScale };

        float in_y = ${ inY };
        if( in_y < 0.0 || in_y > ${ inputHeightFloat } ) {
          setOutput(float(${ extrapolationValue }));
          return;
        }
        float in_x = ${ inX };
        if( in_x < 0.0 || in_x > ${ inputWidthFloat } ) {
          setOutput(float(${ extrapolationValue }));
          return;
        }

        vec2 sourceFracIndexCR = vec2(in_x,in_y);
        if(${ methodId } == 1) {
          // Compute the four integer indices.
          ivec2 sourceFloorCR = ivec2(sourceFracIndexCR);
          ivec2 sourceCeilCR = ivec2(ceil(sourceFracIndexCR));

          float topLeft = getImage(b, sourceFloorCR.y, sourceFloorCR.x, d);
          float bottomLeft = getImage(b, sourceCeilCR.y, sourceFloorCR.x, d);
          float topRight = getImage(b, sourceFloorCR.y, sourceCeilCR.x, d);
          float bottomRight = getImage(b, sourceCeilCR.y, sourceCeilCR.x, d);

          vec2 fracCR = sourceFracIndexCR - vec2(sourceFloorCR);

          float top = topLeft + (topRight - topLeft) * fracCR.x;
          float bottom = bottomLeft + (bottomRight - bottomLeft) * fracCR.x;
          float newValue = top + (bottom - top) * fracCR.y;
          setOutput(newValue);
        } else {
          // Compute the coordinators of nearest neighbor point.
          ivec2 sourceNearestCR = ivec2(floor(
            sourceFracIndexCR + vec2(0.5,0.5)));
          float newValue = getImage(b, sourceNearestCR.y, sourceNearestCR.x, d);
          setOutput(newValue);
        }
      }
    ` ; } }