* Calculate class activation map (CAM) and overlay it on input image. * * This function automatically finds the last convolutional layer, get its * output (activation) under the input image, weights its filters by the * gradient of the class output with respect to them, and then collapses along
(model, classIndex, x, overlayFactor = 2.0)
| 47 | * shape `[1, height, width, 3]`. |
| 48 | */ |
| 49 | function gradClassActivationMap(model, classIndex, x, overlayFactor = 2.0) { |
| 50 | // Try to locate the last conv layer of the model. |
| 51 | let layerIndex = model.layers.length - 1; |
| 52 | while (layerIndex >= 0) { |
| 53 | if (model.layers[layerIndex].getClassName().startsWith('Conv')) { |
| 54 | break; |
| 55 | } |
| 56 | layerIndex--; |
| 57 | } |
| 58 | tf.util.assert( |
| 59 | layerIndex >= 0, `Failed to find a convolutional layer in model`); |
| 60 | |
| 61 | const lastConvLayer = model.layers[layerIndex]; |
| 62 | console.log( |
| 63 | `Located last convolutional layer of the model at ` + |
| 64 | `index ${layerIndex}: layer type = ${lastConvLayer.getClassName()}; ` + |
| 65 | `layer name = ${lastConvLayer.name}`); |
| 66 | |
| 67 | // Get "sub-model 1", which goes from the original input to the output |
| 68 | // of the last convolutional layer. |
| 69 | const lastConvLayerOutput = lastConvLayer.output; |
| 70 | const subModel1 = |
| 71 | tf.model({inputs: model.inputs, outputs: lastConvLayerOutput}); |
| 72 | |
| 73 | // Get "sub-model 2", which goes from the output of the last convolutional |
| 74 | // layer to the original output. |
| 75 | const newInput = tf.input({shape: lastConvLayerOutput.shape.slice(1)}); |
| 76 | layerIndex++; |
| 77 | let y = newInput; |
| 78 | while (layerIndex < model.layers.length) { |
| 79 | y = model.layers[layerIndex++].apply(y); |
| 80 | } |
| 81 | const subModel2 = tf.model({inputs: newInput, outputs: y}); |
| 82 | |
| 83 | return tf.tidy(() => { |
| 84 | // This function runs sub-model 2 and extracts the slice of the probability |
| 85 | // output that corresponds to the desired class. |
| 86 | const convOutput2ClassOutput = (input) => |
| 87 | subModel2.apply(input, {training: true}).gather([classIndex], 1); |
| 88 | // This is the gradient function of the output corresponding to the desired |
| 89 | // class with respect to its input (i.e., the output of the last |
| 90 | // convolutional layer of the original model). |
| 91 | const gradFunction = tf.grad(convOutput2ClassOutput); |
| 92 | |
| 93 | // Calculate the values of the last conv layer's output. |
| 94 | const lastConvLayerOutputValues = subModel1.apply(x); |
| 95 | // Calculate the values of gradients of the class output w.r.t. the output |
| 96 | // of the last convolutional layer. |
| 97 | const gradValues = gradFunction(lastConvLayerOutputValues); |
| 98 | |
| 99 | // Pool the gradient values within each filter of the last convolutional |
| 100 | // layer, resulting in a tensor of shape [numFilters]. |
| 101 | const pooledGradValues = tf.mean(gradValues, [0, 1, 2]); |
| 102 | // Scale the convlutional layer's output by the pooled gradients, using |
| 103 | // broadcasting. |
| 104 | const scaledConvOutputValues = |
| 105 | lastConvLayerOutputValues.mul(pooledGradValues); |
| 106 |