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