async function run()

in 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');
    await utils.writeImageTensorToFile(xWithCAMOverlay, camImagePath);
    console.log(`Written CAM-overlaid image to: ${camImagePath}`);

    // Create manifest and write it to disk.
    const manifest = {
      indices,
      origImagePath,
      classNames,
      probScores,
      layerName2FilePaths,
      layerName2ImageDims,
      camImagePath,
      topIndex: indices[0],
      topProb: probScores[0],
      topClass: classNames[0]
    };
    const manifestPath = path.join(args.outputDir, 'activation-manifest.json');
    fs.writeFileSync(manifestPath, JSON.stringify(manifest));
  } else {
    // Calculate the maximally-activating input images for the conv layers'
    // filters.
    const layerNames = args.convLayerNames.split(',');
    const manifest = {layers: []};
    for (let i = 0; i < layerNames.length; ++i) {
      const layerName = layerNames[i];
      console.log(
          `\n=== Processing layer ${i + 1} of ${layerNames.length}: ` +
          `${layerName} ===`);
      const filePaths = await writeConvLayerFilters(
          model, layerName, args.filters, args.iterations, args.outputDir);
      manifest.layers.push({layerName, filePaths});
    }
    // Write manifest to file.
    const manifestPath = path.join(args.outputDir, 'filters-manifest.json');
    fs.writeFileSync(manifestPath, JSON.stringify(manifest));
  }
};