Function run

visualize-convnet/main.js:109–203

()

}

async function run() {
  const args = parseArguments();
  if (args.gpu) {
    // Use GPU bindings.
    require('@tensorflow/tfjs-node-gpu');
  } else {
    // Use CPU bindings.
    require('@tensorflow/tfjs-node');
  }

  console.log('Loading model...');
  if (args.modelJsonUrl.indexOf('http://') === -1 &&
      args.modelJsonUrl.indexOf('https://') === -1 &&
      args.modelJsonUrl.indexOf('file://') === -1) {
    args.modelJsonUrl = `file://${args.modelJsonUrl}`;
  }
  const model = await tf.loadLayersModel(args.modelJsonUrl);
  console.log('Model loading complete.');

  if (!fs.existsSync(args.outputDir)) {
    shelljs.mkdir('-p', args.outputDir);
  }

  if (args.inputImage != null && args.inputImage !== '') {
    // Compute the internal activations of the conv layers' outputs.
    const imageHeight = model.inputs[0].shape[1];
    const imageWidth = model.inputs[0].shape[2];
    const x = await utils.readImageTensorFromFile(
        args.inputImage, imageHeight, imageWidth);
    const layerNames = args.convLayerNames.split(',');
    const {modelOutput, layerName2FilePaths, layerName2ImageDims} =
        await filters.writeInternalActivationAndGetOutput(
            model, layerNames, x, args.filters, args.outputDir);

    // Calculate internal activations and final output of the model.
    const topNum = 10;
    const {values: topKVals, indices: topKIndices} =
        tf.topk(modelOutput, topNum);
    const probScores = Array.from(await topKVals.data());
    const indices = Array.from(await topKIndices.data());
    const classNames =
        indices.map(index => imagenetClasses.IMAGENET_CLASSES[index]);

    console.log(`Top-${topNum} classes:`);
    for (let i = 0; i < topNum; ++i) {
      console.log(
          `  ${classNames[i]} (index=${indices[i]}): ` +
          `${probScores[i].toFixed(4)}`);
    }

    // Save the original input image and the top-10 classification results.
    const origImagePath =
        path.join(args.outputDir, path.basename(args.inputImage));
    shelljs.cp(args.inputImage, origImagePath);

    // Calculate Grad-CAM heatmap.
    const xWithCAMOverlay = cam.gradClassActivationMap(model, indices[0], x);
    const camImagePath = path.join(args.outputDir, 'cam.png');