Add support for converting models with ImageProjectiveTransformV3 op (#…

…6206) * Add support for converting models with ImageProjectiveTransformV3 op * Rename image to images Co-authored-by: Na Li <linazhao@google.com>
tensorflow · Mar 15, 2022 · 00d949a · 00d949a
1 parent 092c976
commit 00d949a
Show file tree

Hide file tree

Showing 4 changed files with 98 additions and 1 deletion.
diff --git a/tfjs-converter/docs/supported_ops.md b/tfjs-converter/docs/supported_ops.md
@@ -230,6 +230,7 @@
 |ResizeNearestNeighbor|resizeNearestNeighbor|
 |Not mapped|flipLeftRight|
 |Not mapped|rotateWithOffset|
+|ImageProjectiveTransformV3|transform|
 
 ## Operations - Matrices
 

diff --git a/tfjs-converter/python/tensorflowjs/op_list/image.json b/tfjs-converter/python/tensorflowjs/op_list/image.json
@@ -104,5 +104,43 @@
         "type": "number"
       }
     ]
+  },
+  {
+    "tfOpName": "ImageProjectiveTransformV3",
+    "category": "image",
+    "inputs": [
+      {
+        "start": 0,
+        "name": "images",
+        "type": "tensor"
+      },
+      {
+        "start": 1,
+        "name": "transforms",
+        "type": "tensor"
+      },
+      {
+        "start": 2,
+        "name": "outputShape",
+        "type": "number[]"
+      },
+      {
+        "start": 3,
+        "name": "fillValue",
+        "type": "number"
+      }
+    ],
+    "attrs": [
+      {
+        "tfName": "interpolation",
+        "name": "interpolation",
+        "type": "string"
+      },
+      {
+        "tfName": "fill_mode",
+        "name": "fillMode",
+        "type": "string"
+      }
+    ]
   }
 ]
diff --git a/tfjs-converter/src/operations/executors/image_executor.ts b/tfjs-converter/src/operations/executors/image_executor.ts
@@ -78,6 +78,29 @@ export const executeOp: InternalOpExecutor =
               cropSize as [number, number], method as 'bilinear' | 'nearest',
               extrapolationValue)];
         }
+        case 'ImageProjectiveTransformV3': {
+          const images =
+              getParamValue('images', node, tensorMap, context) as Tensor;
+          const transforms =
+              getParamValue('transforms', node, tensorMap, context) as Tensor;
+          const outputShape =
+              getParamValue('outputShape', node, tensorMap, context) as
+              number[];
+          const fillValue =
+              getParamValue('fillValue', node, tensorMap, context) as number;
+          const interpolation =
+              getParamValue('interpolation', node, tensorMap, context) as
+              string;
+          const fillMode =
+              getParamValue('fillMode', node, tensorMap, context) as string;
+          return [tfOps.image.transform(
+              images as Tensor4D,
+              transforms as Tensor2D,
+              interpolation.toLowerCase() as 'bilinear' | 'nearest',
+              fillMode.toLowerCase() as 'constant' | 'reflect' | 'wrap' | 'nearest',
+              fillValue,
+              outputShape as [number, number])];
+        }
         default:
           throw TypeError(`Node type ${node.op} is not implemented`);
       }

diff --git a/tfjs-converter/src/operations/executors/image_executor_test.ts b/tfjs-converter/src/operations/executors/image_executor_test.ts
@@ -22,7 +22,7 @@ import * as image from '../op_list/image';
 import {Node} from '../types';
 
 import {executeOp} from './image_executor';
-import {createBoolAttr, createNumberAttr, createNumericArrayAttrFromIndex, createStrAttr, createTensorAttr, validateParam} from './test_helper';
+import {createBoolAttr, createNumberAttr, createNumberAttrFromIndex, createNumericArrayAttrFromIndex, createStrAttr, createTensorAttr, validateParam} from './test_helper';
 
 describe('image', () => {
   let node: Node;
@@ -123,6 +123,41 @@ describe('image', () => {
         node.attrParams['extrapolationValue'] = createNumberAttr(0.5);
         node.inputNames = ['input1', 'input2', 'input3', 'input4'];
 
+        expect(validateParam(node, image.json)).toBeTruthy();
+      });
+    });
+    describe('ImageProjectiveTransformV3', () => {
+      it('should return input', () => {
+        node.op = 'ImageProjectiveTransformV3';
+        node.inputParams['images'] = createTensorAttr(0);
+        node.inputParams['transforms'] = createTensorAttr(1);
+        node.inputParams['outputShape'] = createNumericArrayAttrFromIndex(2);
+        node.inputParams['fillValue'] = createNumberAttrFromIndex(3);
+        node.attrParams['interpolation'] = createStrAttr('bilinear');
+        node.attrParams['fillMode'] = createStrAttr('constant');
+        node.inputNames = ['input1', 'input2', 'input3', 'input4'];
+
+        spyOn(tfOps.image, 'transform');
+        const input2 = [tfOps.tensor1d([2])];
+        const input3 = [tfOps.tensor1d([4, 5])];
+        const input4 = [tfOps.scalar(3)];
+
+        executeOp(node, {input1, input2, input3, input4}, context);
+        expect(tfOps.image.transform)
+            .toHaveBeenCalledWith(
+                input1[0], input2[0], 'bilinear', 'constant', 3, [4, 5]);
+      });
+
+      it('should match json def', () => {
+        node.op = 'ImageProjectiveTransformV3';
+        node.inputParams['images'] = createTensorAttr(0);
+        node.inputParams['transforms'] = createTensorAttr(1);
+        node.inputParams['outputShape'] = createNumericArrayAttrFromIndex(2);
+        node.inputParams['fillValue'] = createNumberAttrFromIndex(3);
+        node.attrParams['interpolation'] = createStrAttr('bilinear');
+        node.attrParams['fillMode'] = createStrAttr('constant');
+        node.inputNames = ['input1', 'input2', 'input3', 'input4'];
+
         expect(validateParam(node, image.json)).toBeTruthy();
       });
     });