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); }
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); }
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; });
residualBlock(input, id) { const conv1 = this.convLayer(input, 1, true, id); const conv2 = this.convLayer(conv1, 1, false, id + 3); return tf.add(conv2, input); }