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));
}
};