Commit 084cd30a by vincent

some fixes in stage1

parent 9ffe1cc9
...@@ -4,6 +4,7 @@ import { Point } from '../Point'; ...@@ -4,6 +4,7 @@ import { Point } from '../Point';
import { BoundingBox } from './BoundingBox'; import { BoundingBox } from './BoundingBox';
import { CELL_SIZE, CELL_STRIDE } from './config'; import { CELL_SIZE, CELL_STRIDE } from './config';
import { nms } from './nms'; import { nms } from './nms';
import { normalize } from './normalize';
import { PNet } from './PNet'; import { PNet } from './PNet';
import { PNetParams } from './types'; import { PNetParams } from './types';
...@@ -12,13 +13,12 @@ function rescaleAndNormalize(x: tf.Tensor4D, scale: number): tf.Tensor4D { ...@@ -12,13 +13,12 @@ function rescaleAndNormalize(x: tf.Tensor4D, scale: number): tf.Tensor4D {
const [height, width] = x.shape.slice(1) const [height, width] = x.shape.slice(1)
const resized = tf.image.resizeBilinear(x, [Math.floor(height * scale), Math.floor(width * scale)]) const resized = tf.image.resizeBilinear(x, [Math.floor(height * scale), Math.floor(width * scale)])
const normalized = tf.mul(tf.sub(resized, tf.scalar(127.5)), tf.scalar(0.0078125)) const normalized = normalize(resized)
return (tf.transpose(normalized, [0, 2, 1, 3]) as tf.Tensor4D) return (tf.transpose(normalized, [0, 2, 1, 3]) as tf.Tensor4D)
}) })
} }
function extractBoundingBoxes( function extractBoundingBoxes(
scoresTensor: tf.Tensor2D, scoresTensor: tf.Tensor2D,
regionsTensor: tf.Tensor3D, regionsTensor: tf.Tensor3D,
...@@ -38,10 +38,10 @@ function extractBoundingBoxes( ...@@ -38,10 +38,10 @@ function extractBoundingBoxes(
const boundingBoxes = indices.map(idx => { const boundingBoxes = indices.map(idx => {
const cell = new BoundingBox( const cell = new BoundingBox(
Math.round((idx.x * CELL_STRIDE + 1) / scale),
Math.round((idx.y * CELL_STRIDE + 1) / scale), Math.round((idx.y * CELL_STRIDE + 1) / scale),
Math.round((idx.x * CELL_STRIDE + CELL_SIZE) / scale), Math.round((idx.x * CELL_STRIDE + 1) / scale),
Math.round((idx.y * CELL_STRIDE + CELL_SIZE) / scale) Math.round((idx.y * CELL_STRIDE + CELL_SIZE) / scale),
Math.round((idx.x * CELL_STRIDE + CELL_SIZE) / scale)
) )
const score = scoresTensor.get(idx.y, idx.x) const score = scoresTensor.get(idx.y, idx.x)
...@@ -63,14 +63,22 @@ function extractBoundingBoxes( ...@@ -63,14 +63,22 @@ function extractBoundingBoxes(
return boundingBoxes return boundingBoxes
} }
export function stage1(imgTensor: tf.Tensor4D, scales: number[], scoreThreshold: number, params: PNetParams) { export function stage1(
imgTensor: tf.Tensor4D,
scales: number[],
scoreThreshold: number,
params: PNetParams
) {
const boxesForScale = scales.map((scale) => { const boxesForScale = scales.map((scale, i) => {
const { scoresTensor, regionsTensor } = tf.tidy(() => { const { scoresTensor, regionsTensor } = tf.tidy(() => {
const resized = rescaleAndNormalize(imgTensor, scale) const resized = rescaleAndNormalize(imgTensor, scale)
const { prob, regions } = PNet(resized, params) const { prob, regions } = PNet(resized, params)
const scores = tf.unstack(prob, 3)[1] const scores = tf.unstack(prob, 3)[1]
const [sh, sw] = scores.shape.slice(1) const [sh, sw] = scores.shape.slice(1)
const [rh, rw] = regions.shape.slice(1) const [rh, rw] = regions.shape.slice(1)
...@@ -120,9 +128,9 @@ export function stage1(imgTensor: tf.Tensor4D, scales: number[], scoreThreshold: ...@@ -120,9 +128,9 @@ export function stage1(imgTensor: tf.Tensor4D, scales: number[], scoreThreshold:
.map(({ cell, region, score }) => ({ .map(({ cell, region, score }) => ({
box: new BoundingBox( box: new BoundingBox(
cell.left + (region.left * cell.width), cell.left + (region.left * cell.width),
cell.right + (region.right * cell.width),
cell.top + (region.top * cell.height), cell.top + (region.top * cell.height),
cell.bottom + (region.bottom * cell.height), cell.right + (region.right * cell.width),
cell.bottom + (region.bottom * cell.height)
).toSquare().round(), ).toSquare().round(),
score score
})) }))
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment