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 function yolo_boxes_to_corners(box_xy, box_wh) {
const two = tf.tensor1d([2.0]);
const box_mins = tf.sub(box_xy, tf.div(box_wh, two));
const box_maxes = tf.add(box_xy, tf.div(box_wh, two));
const dim_0 = box_mins.shape[0];
const dim_1 = box_mins.shape[1];
const dim_2 = box_mins.shape[2];
const size = [dim_0, dim_1, dim_2, 1];
return tf.concat([
box_mins.slice([0, 0, 0, 1], size),
box_mins.slice([0, 0, 0, 0], size),
box_maxes.slice([0, 0, 0, 1], size),
box_maxes.slice([0, 0, 0, 0], size),
], 3);
}
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;
});