outputLayer.js 2.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950
  1. "use strict";
  2. Object.defineProperty(exports, "__esModule", { value: true });
  3. var tf = require("@tensorflow/tfjs-core");
  4. function getCenterCoordinatesAndSizesLayer(x) {
  5. var vec = tf.unstack(tf.transpose(x, [1, 0]));
  6. var sizes = [
  7. tf.sub(vec[2], vec[0]),
  8. tf.sub(vec[3], vec[1])
  9. ];
  10. var centers = [
  11. tf.add(vec[0], tf.div(sizes[0], tf.scalar(2))),
  12. tf.add(vec[1], tf.div(sizes[1], tf.scalar(2)))
  13. ];
  14. return {
  15. sizes: sizes,
  16. centers: centers
  17. };
  18. }
  19. function decodeBoxesLayer(x0, x1) {
  20. var _a = getCenterCoordinatesAndSizesLayer(x0), sizes = _a.sizes, centers = _a.centers;
  21. var vec = tf.unstack(tf.transpose(x1, [1, 0]));
  22. var div0_out = tf.div(tf.mul(tf.exp(tf.div(vec[2], tf.scalar(5))), sizes[0]), tf.scalar(2));
  23. var add0_out = tf.add(tf.mul(tf.div(vec[0], tf.scalar(10)), sizes[0]), centers[0]);
  24. var div1_out = tf.div(tf.mul(tf.exp(tf.div(vec[3], tf.scalar(5))), sizes[1]), tf.scalar(2));
  25. var add1_out = tf.add(tf.mul(tf.div(vec[1], tf.scalar(10)), sizes[1]), centers[1]);
  26. return tf.transpose(tf.stack([
  27. tf.sub(add0_out, div0_out),
  28. tf.sub(add1_out, div1_out),
  29. tf.add(add0_out, div0_out),
  30. tf.add(add1_out, div1_out)
  31. ]), [1, 0]);
  32. }
  33. function outputLayer(boxPredictions, classPredictions, params) {
  34. return tf.tidy(function () {
  35. var batchSize = boxPredictions.shape[0];
  36. var boxes = decodeBoxesLayer(tf.reshape(tf.tile(params.extra_dim, [batchSize, 1, 1]), [-1, 4]), tf.reshape(boxPredictions, [-1, 4]));
  37. boxes = tf.reshape(boxes, [batchSize, (boxes.shape[0] / batchSize), 4]);
  38. var scoresAndClasses = tf.sigmoid(tf.slice(classPredictions, [0, 0, 1], [-1, -1, -1]));
  39. var scores = tf.slice(scoresAndClasses, [0, 0, 0], [-1, -1, 1]);
  40. scores = tf.reshape(scores, [batchSize, scores.shape[1]]);
  41. var boxesByBatch = tf.unstack(boxes);
  42. var scoresByBatch = tf.unstack(scores);
  43. return {
  44. boxes: boxesByBatch,
  45. scores: scoresByBatch
  46. };
  47. });
  48. }
  49. exports.outputLayer = outputLayer;
  50. //# sourceMappingURL=outputLayer.js.map