export function yolo_head(feats, anchors, num_classes) { const num_anchors = anchors.shape[0]; const anchors_tensor = tf.reshape(anchors, [1, 1, num_anchors, 2]); let conv_dims = feats.shape.slice(1, 3); // For later use const conv_dims_0 = conv_dims[0]; const conv_dims_1 = conv_dims[1]; let conv_height_index = tf.range(0, conv_dims[0]); let conv_width_index = tf.range(0, conv_dims[1]); conv_height_index = tf.tile(conv_height_index, [conv_dims[1]]) conv_width_index = tf.tile(tf.expandDims(conv_width_index, 0), [conv_dims[0], 1]); conv_width_index = tf.transpose(conv_width_index).flatten(); let conv_index = tf.transpose(tf.stack([conv_height_index, conv_width_index])); conv_index = tf.reshape(conv_index, [conv_dims[0], conv_dims[1], 1, 2]) conv_index = tf.cast(conv_index, feats.dtype); feats = tf.reshape(feats, [conv_dims[0], conv_dims[1], num_anchors, num_classes + 5]); conv_dims = tf.cast(tf.reshape(tf.tensor1d(conv_dims), [1,1,1,2]), feats.dtype); let box_xy = tf.sigmoid(feats.slice([0,0,0,0], [conv_dims_0, conv_dims_1, num_anchors, 2])) let box_wh = tf.exp(feats.slice([0,0,0, 2], [conv_dims_0, conv_dims_1, num_anchors, 2])) const box_confidence = tf.sigmoid(feats.slice([0,0,0, 4], [conv_dims_0, conv_dims_1, num_anchors, 1])) const box_class_probs = tf.softmax(feats.slice([0,0,0, 5],[conv_dims_0, conv_dims_1, num_anchors, num_classes])); box_xy = tf.div(tf.add(box_xy, conv_index), conv_dims); box_wh = tf.div(tf.mul(box_wh, anchors_tensor), conv_dims); return [ box_xy, box_wh, box_confidence, box_class_probs ]; }
export async function yolo_filter_boxes( boxes, box_confidence, box_class_probs, threshold ) { const box_scores = tf.mul(box_confidence, box_class_probs); const box_classes = tf.argMax(box_scores, -1); const box_class_scores = tf.max(box_scores, -1); const prediction_mask = tf.greaterEqual(box_class_scores, tf.scalar(threshold)); const mask_arr = await prediction_mask.data(); const indices_arr = []; for (let i=0; i<mask_arr.length; i++) { const v = mask_arr[i]; if (v) { indices_arr.push(i); } } if (indices_arr.length == 0) { return [null, null, null]; } const indices = tf.tensor1d(indices_arr); return [ tf.gather(boxes.reshape([mask_arr.length, 4]), indices), tf.gather(box_class_scores.flatten(), indices), tf.gather(box_classes.flatten(), indices), ]; }
instanceNorm(input, id) { const [height, width, inDepth] = input.shape; const moments = tf.moments(input, [0, 1]); const mu = moments.mean; const sigmaSq = moments.variance; const shift = this.variables[StyleTransfer.getVariableName(id)]; const scale = this.variables[StyleTransfer.getVariableName(id + 1)]; const epsilon = this.epsilonScalar; const normalized = tf.div(tf.sub(input.asType('float32'), mu), tf.sqrt(tf.add(sigmaSq, epsilon))); const shifted = tf.add(tf.mul(scale, normalized), shift); return shifted.as3D(height, width, inDepth); }
const result = tf.tidy(() => { const conv1 = this.convLayer(image, 1, true, 0); const conv2 = this.convLayer(conv1, 2, true, 3); const conv3 = this.convLayer(conv2, 2, true, 6); const res1 = this.residualBlock(conv3, 9); const res2 = this.residualBlock(res1, 15); const res3 = this.residualBlock(res2, 21); const res4 = this.residualBlock(res3, 27); const res5 = this.residualBlock(res4, 33); const convT1 = this.convTransposeLayer(res5, 64, 2, 39); const convT2 = this.convTransposeLayer(convT1, 32, 2, 42); const convT3 = this.convLayer(convT2, 1, false, 45); const outTanh = tf.tanh(convT3); const scaled = tf.mul(this.timesScalar, outTanh); const shifted = tf.add(this.plusScalar, scaled); const clamped = tf.clipByValue(shifted, 0, 255); const normalized = tf.div(clamped, tf.scalar(255.0)); return normalized; });