Skip to content

Commit

Permalink
[tfjs-node] Encode jpeg (tensorflow#7484)
Browse files Browse the repository at this point in the history 
fixed tensorflow#7459
ensure TF tensors are disposed after encode image op is called.
  • Loading branch information
@pyu10055
pyu10055 committed Mar 15, 2023
1 parent 86e7f4a commit e1fd1f1
Show file tree
Hide file tree
Showing 6 changed files with 110 additions and 3 deletions.
9 changes: 9 additions & 0 deletions tfjs-node/binding/tfjs_backend.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -804,6 +804,15 @@ void TFJSBackend::DeleteTensor(napi_env env, napi_value tensor_id_value) {
tfe_handle_map_.erase(tensor_entry);
}

napi_value TFJSBackend::GetNumOfTensors(napi_env env) {
napi_status nstatus;
napi_value num_of_tensors;
nstatus =
napi_create_int32(env, tfe_handle_map_.size(), &num_of_tensors);
ENSURE_NAPI_OK_RETVAL(env, nstatus, nullptr);
return num_of_tensors;
}

napi_value TFJSBackend::GetTensorData(napi_env env,
napi_value tensor_id_value) {
int32_t tensor_id;
Expand Down
3 changes: 3 additions & 0 deletions tfjs-node/binding/tfjs_backend.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,9 @@ class TFJSBackend {
// Get number of loaded SavedModel in the backend:
napi_value GetNumOfSavedModels(napi_env env);

// Get number of tensor handles in the backend:
napi_value GetNumOfTensors(napi_env env);

private:
TFJSBackend(napi_env env);

Expand Down
10 changes: 10 additions & 0 deletions tfjs-node/binding/tfjs_binding.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,13 @@ static napi_value RunSavedModel(napi_env env, napi_callback_info info) {
return backend->RunSavedModel(env, args[0], args[1], args[2], args[3]);
}

static napi_value GetNumOfTensors(napi_env env, napi_callback_info info) {
TFJSBackend *const backend = GetTFJSBackend(env);
if (!backend) return nullptr;
// Delete SavedModel takes 0 param;
return backend->GetNumOfTensors(env);
}

static napi_value GetNumOfSavedModels(napi_env env, napi_callback_info info) {
TFJSBackend *const backend = GetTFJSBackend(env);
if (!backend) return nullptr;
Expand Down Expand Up @@ -314,6 +321,9 @@ static napi_value InitTFNodeJSBinding(napi_env env, napi_value exports) {
napi_default, nullptr},
{"getNumOfSavedModels", nullptr, GetNumOfSavedModels, nullptr, nullptr,
nullptr, napi_default, nullptr},
{"getNumOfTensors", nullptr, GetNumOfTensors, nullptr, nullptr,
nullptr, napi_default, nullptr},

};
nstatus = napi_define_properties(env, exports, ARRAY_SIZE(exports_properties),
exports_properties);
Expand Down
80 changes: 80 additions & 0 deletions tfjs-node/src/image_test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,15 @@ import {promisify} from 'util';

import {getImageType, ImageType} from './image';
import * as tf from './index';
import {NodeJSKernelBackend} from './nodejs_kernel_backend';

const readFile = promisify(fs.readFile);

describe('decode images', () => {
it('decode png', async () => {
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const uint8array =
await getUint8ArrayFromImage('test_objects/images/image_png_test.png');
const imageTensor = tf.node.decodePng(uint8array);
Expand All @@ -38,21 +41,29 @@ describe('decode images', () => {
await imageTensor.data(),
[238, 101, 0, 50, 50, 50, 100, 50, 0, 200, 100, 50]);
expect(memory().numTensors).toBe(beforeNumTensors + 1);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors + 1);
});

it('decode png 1 channels', async () => {
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const uint8array =
await getUint8ArrayFromImage('test_objects/images/image_png_test.png');
const imageTensor = tf.node.decodeImage(uint8array, 1);
expect(imageTensor.dtype).toBe('int32');
expect(imageTensor.shape).toEqual([2, 2, 1]);
test_util.expectArraysEqual(await imageTensor.data(), [130, 50, 59, 124]);
expect(memory().numTensors).toBe(beforeNumTensors + 1);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors + 1);
});

it('decode png 3 channels', async () => {
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const uint8array =
await getUint8ArrayFromImage('test_objects/images/image_png_test.png');
const imageTensor = tf.node.decodeImage(uint8array);
Expand All @@ -62,10 +73,14 @@ describe('decode images', () => {
await imageTensor.data(),
[238, 101, 0, 50, 50, 50, 100, 50, 0, 200, 100, 50]);
expect(memory().numTensors).toBe(beforeNumTensors + 1);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors + 1);
});

it('decode png 4 channels', async () => {
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const uint8array = await getUint8ArrayFromImage(
'test_objects/images/image_png_4_channel_test.png');
const imageTensor = tf.node.decodeImage(uint8array, 4);
Expand All @@ -75,10 +90,14 @@ describe('decode images', () => {
238, 101, 0, 255, 50, 50, 50, 255, 100, 50, 0, 255, 200, 100, 50, 255
]);
expect(memory().numTensors).toBe(beforeNumTensors + 1);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors + 1);
});

it('decode bmp', async () => {
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const uint8array =
await getUint8ArrayFromImage('test_objects/images/image_bmp_test.bmp');
const imageTensor = tf.node.decodeBmp(uint8array);
Expand All @@ -88,10 +107,14 @@ describe('decode images', () => {
await imageTensor.data(),
[238, 101, 0, 50, 50, 50, 100, 50, 0, 200, 100, 50]);
expect(memory().numTensors).toBe(beforeNumTensors + 1);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors + 1);
});

it('decode bmp through decodeImage', async () => {
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const uint8array =
await getUint8ArrayFromImage('test_objects/images/image_bmp_test.bmp');
const imageTensor = tf.node.decodeImage(uint8array);
Expand All @@ -101,10 +124,15 @@ describe('decode images', () => {
await imageTensor.data(),
[238, 101, 0, 50, 50, 50, 100, 50, 0, 200, 100, 50]);
expect(memory().numTensors).toBe(beforeNumTensors + 1);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors + 1);
});

it('decode jpg', async () => {
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();

const uint8array = await getUint8ArrayFromImage(
'test_objects/images/image_jpeg_test.jpeg');
const imageTensor = tf.node.decodeJpeg(uint8array);
Expand All @@ -114,6 +142,8 @@ describe('decode images', () => {
await imageTensor.data(),
[239, 100, 0, 46, 48, 47, 92, 49, 0, 194, 98, 47]);
expect(memory().numTensors).toBe(beforeNumTensors + 1);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors + 1);
});

it('decode jpeg node bindings do not leak', async () => {
Expand Down Expand Up @@ -142,17 +172,24 @@ describe('decode images', () => {

it('decode jpg 1 channel', async () => {
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();

const uint8array = await getUint8ArrayFromImage(
'test_objects/images/image_jpeg_test.jpeg');
const imageTensor = tf.node.decodeImage(uint8array, 1);
expect(imageTensor.dtype).toBe('int32');
expect(imageTensor.shape).toEqual([2, 2, 1]);
test_util.expectArraysEqual(await imageTensor.data(), [130, 47, 56, 121]);
expect(memory().numTensors).toBe(beforeNumTensors + 1);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors + 1);
});

it('decode jpg 3 channels', async () => {
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const uint8array = await getUint8ArrayFromImage(
'test_objects/images/image_jpeg_test.jpeg');
const imageTensor = tf.node.decodeImage(uint8array, 3);
Expand All @@ -162,12 +199,16 @@ describe('decode images', () => {
await imageTensor.data(),
[239, 100, 0, 46, 48, 47, 92, 49, 0, 194, 98, 47]);
expect(memory().numTensors).toBe(beforeNumTensors + 1);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors + 1);
});

it('decode jpg with 0 channels, use the number of channels in the ' +
'JPEG-encoded image',
async () => {
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const uint8array = await getUint8ArrayFromImage(
'test_objects/images/image_jpeg_test.jpeg');
const imageTensor = tf.node.decodeImage(uint8array);
Expand All @@ -177,21 +218,29 @@ describe('decode images', () => {
await imageTensor.data(),
[239, 100, 0, 46, 48, 47, 92, 49, 0, 194, 98, 47]);
expect(memory().numTensors).toBe(beforeNumTensors + 1);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors + 1);
});

it('decode jpg with downscale', async () => {
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const uint8array = await getUint8ArrayFromImage(
'test_objects/images/image_jpeg_test.jpeg');
const imageTensor = tf.node.decodeJpeg(uint8array, 0, 2);
expect(imageTensor.dtype).toBe('int32');
expect(imageTensor.shape).toEqual([1, 1, 3]);
test_util.expectArraysEqual(await imageTensor.data(), [147, 75, 25]);
expect(memory().numTensors).toBe(beforeNumTensors + 1);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors + 1);
});

it('decode gif', async () => {
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const uint8array =
await getUint8ArrayFromImage('test_objects/images/gif_test.gif');
const imageTensor = tf.node.decodeImage(uint8array);
Expand All @@ -202,10 +251,15 @@ describe('decode images', () => {
200, 100, 50, 34, 68, 102, 170, 0, 102, 255, 255, 255
]);
expect(memory().numTensors).toBe(beforeNumTensors + 1);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors + 1);
});

it('decode gif with no expandAnimation', async () => {
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();

const uint8array =
await getUint8ArrayFromImage('test_objects/images/gif_test.gif');
const imageTensor = tf.node.decodeImage(uint8array, 3, 'int32', false);
Expand All @@ -215,6 +269,8 @@ describe('decode images', () => {
await imageTensor.data(),
[238, 101, 0, 50, 50, 50, 100, 50, 0, 200, 100, 50]);
expect(memory().numTensors).toBe(beforeNumTensors + 1);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors + 1);
});

it('throw error if request non int32 dtype', async done => {
Expand Down Expand Up @@ -271,18 +327,26 @@ describe('encode images', () => {
new Uint8Array([239, 100, 0, 46, 48, 47, 92, 49, 0, 194, 98, 47]),
[2, 2, 3]);
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const jpegEncodedData = await tf.node.encodeJpeg(imageTensor);
expect(memory().numTensors).toBe(beforeNumTensors);
expect(getImageType(jpegEncodedData)).toEqual(ImageType.JPEG);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors);
imageTensor.dispose();
});

it('encodeJpeg grayscale', async () => {
const imageTensor = tf.tensor3d(new Uint8Array([239, 0, 47, 0]), [2, 2, 1]);
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const jpegEncodedData = await tf.node.encodeJpeg(imageTensor, 'grayscale');
expect(memory().numTensors).toBe(beforeNumTensors);
expect(getImageType(jpegEncodedData)).toEqual(ImageType.JPEG);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors);
imageTensor.dispose();
});

Expand All @@ -300,10 +364,14 @@ describe('encode images', () => {
const yDensity = 500;

const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const jpegEncodedData = await tf.node.encodeJpeg(
imageTensor, format, quality, progressive, optimizeSize,
chromaDownsampling, densityUnit, xDensity, yDensity);
expect(memory().numTensors).toBe(beforeNumTensors);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors);
expect(getImageType(jpegEncodedData)).toEqual(ImageType.JPEG);
imageTensor.dispose();
});
Expand All @@ -313,11 +381,15 @@ describe('encode images', () => {
new Uint8Array([239, 100, 0, 46, 48, 47, 92, 49, 0, 194, 98, 47]),
[2, 2, 3]);
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const pngEncodedData = await tf.node.encodePng(imageTensor);
const pngDecodedTensor = await tf.node.decodePng(pngEncodedData);
const pngDecodedData = await pngDecodedTensor.data();
pngDecodedTensor.dispose();
expect(memory().numTensors).toBe(beforeNumTensors);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors);
expect(getImageType(pngEncodedData)).toEqual(ImageType.PNG);
test_util.expectArraysEqual(await imageTensor.data(), pngDecodedData);
imageTensor.dispose();
Expand All @@ -326,8 +398,12 @@ describe('encode images', () => {
it('encodePng grayscale', async () => {
const imageTensor = tf.tensor3d(new Uint8Array([239, 0, 47, 0]), [2, 2, 1]);
const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const pngEncodedData = await tf.node.encodePng(imageTensor);
expect(memory().numTensors).toBe(beforeNumTensors);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors);
expect(getImageType(pngEncodedData)).toEqual(ImageType.PNG);
imageTensor.dispose();
});
Expand All @@ -339,8 +415,12 @@ describe('encode images', () => {
const compression = 4;

const beforeNumTensors = memory().numTensors;
const beforeNumTFTensors =
(tf.backend() as NodeJSKernelBackend).getNumOfTFTensors();
const pngEncodedData = await tf.node.encodePng(imageTensor, compression);
expect(memory().numTensors).toBe(beforeNumTensors);
expect((tf.backend() as NodeJSKernelBackend).getNumOfTFTensors())
.toBe(beforeNumTFTensors);
expect(getImageType(pngEncodedData)).toEqual(ImageType.PNG);
imageTensor.dispose();
});
Expand Down
9 changes: 7 additions & 2 deletions tfjs-node/src/nodejs_kernel_backend.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -120,8 +120,8 @@ export class NodeJSKernelBackend extends KernelBackend {
// We can then change the return type from Tensor to TensorInfo.
// return {dataId: newId, shape: metadata.shape, dtype};

const tensorInfo: TensorInfo = {
dataId: newId, shape: metadata.shape, dtype};
const tensorInfo:
TensorInfo = {dataId: newId, shape: metadata.shape, dtype};
return tf.engine().makeTensorFromTensorInfo(tensorInfo);
}

Expand Down Expand Up @@ -393,6 +393,7 @@ export class NodeJSKernelBackend extends KernelBackend {
this.binding.createTensor(imageShape, this.binding.TF_UINT8, imageData);
const outputMetadata =
this.binding.executeOp(name, opAttrs, [inputTensorId], 1);
this.binding.deleteTensor(inputTensorId);
const outputTensorInfo = outputMetadata[0];
// prevent the tensor data from being converted to a UTF8 string, since
// the encoded data is not valid UTF8
Expand Down Expand Up @@ -680,6 +681,10 @@ export class NodeJSKernelBackend extends KernelBackend {
getNumOfSavedModels() {
return this.binding.getNumOfSavedModels();
}

getNumOfTFTensors() {
return this.binding.getNumOfTensors();
}
}

/** Returns an instance of the Node.js backend. */
Expand Down
2 changes: 1 addition & 1 deletion tfjs-node/src/tfjs_binding.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ export interface TFJSBinding {
outputOpNames: string): TensorMetadata[];

getNumOfSavedModels(): number;

getNumOfTensors(): number;
isUsingGpuDevice(): boolean;

// TF Types
Expand Down

0 comments on commit e1fd1f1

Please sign in to comment.