Skip to content
Snippets Groups Projects
shapes.js 5.26 KiB
function dtor(a) {
  return (Math.PI * a) / 180
}
function cos(a) {
  return Math.cos(dtor(a))
}
function sin(a) {
  return Math.sin(dtor(a))
}

const angleStep = 10
const numAngleCells = 180 / angleStep
const rhoMax = 1000

const getDistance = (a, b) => {
  if (!(a && b)) return 0
  return Math.sqrt(
    (a[0] - b[0]) * (a[0] - b[0]) + (a[1] - b[1]) * (a[1] - b[1]),
  )
}

function findMaxInHough(accum, threshold) {
  let max = 0
  //   let bestRho = 0
  let bestTheta = 0
  for (let i = 0; i < numAngleCells; i++) {
    if (!accum[i]) continue
    for (let j = 0; j < accum[i].length; j++) {
      if (accum[i][j] > max) {
        max = accum[i][j]
        // bestRho = j
        bestTheta = i
      }
    }
  }
  //   bestRho <<= 1
  //   bestRho -= rhoMax
  //   bestRho *= rhoStep
  bestTheta *= angleStep

  if (max > threshold) {
    return bestTheta
  }
  return undefined
}

function constructHoughAccumulator(config, accumulator, x, y) {
  for (let thetaIndex = 0; thetaIndex < numAngleCells; thetaIndex++) {
    const theta = thetaIndex * angleStep
    let rho = x * cos(theta) + y * sin(theta)
    rho = Math.floor(rho)
    rho += rhoMax
    rho >>= 1
    rho /= config.rhoStep
    rho = Math.floor(rho)
    if (accumulator[thetaIndex] == undefined) accumulator[thetaIndex] = []
    if (accumulator[thetaIndex][rho] == undefined) {
      accumulator[thetaIndex][rho] = 1
    } else {
      accumulator[thetaIndex][rho]++
    }
  }
}

function boundingCoords(points) {
  const xs = points.map((p) => p[0])
  const ys = points.map((p) => p[1])
  return {
    maxX: Math.max(...xs),
    minX: Math.min(...xs),
    maxY: Math.max(...ys),
    minY: Math.min(...ys),
  }
}

const LINE_Q = 10

function couldBeLine(points) {
  const { maxX, minX, maxY, minY } = boundingCoords(points)
  const [dx, dy] = [maxX - minX, maxY - minY].map((x) => x + 0.00001)
  return dy / dx > LINE_Q || dx / dy > LINE_Q
}

const MATRIX_SIZE = 3
const MATRIX_CENTER_RATIO = 0.65

function mArray(min, max) {
  const d = max - min
  const centerSegmentSize = d * MATRIX_CENTER_RATIO
  const smallStep = (d - centerSegmentSize) / 2
  const p = [min + smallStep, min + smallStep + centerSegmentSize, max]
  return p
}

function getCluster([x, y], xBounds, yBounds) {
  return {
    x: xBounds.findIndex((bound) => x <= bound),
    y: yBounds.findIndex((bound) => y <= bound),
  }
}

function computeClusters(points, xBounds, yBounds) {
  const clusters = Array(MATRIX_SIZE)
    .fill(0)
    .map(() =>
      Array(MATRIX_SIZE)
        .fill()
        .map(() => ({ arr: [], sum: 0 })),
    )
  const intervals = points.map((point, i) => ({
    point,
    dist: getDistance(point, points[i + 1]),
  }))

  let totalSum = 0
  intervals.forEach((interval) => {
    const { x, y } = getCluster(interval.point, xBounds, yBounds)
    clusters[x][y].arr.push(interval)
    clusters[x][y].sum += interval.dist
    totalSum += interval.dist
  })

  return { arr: clusters, totalSum }
}

function clusterCoefficients(clusters) {
  return clusters.arr.map((rowCluster) =>
    rowCluster.map((cluster) => cluster.sum / clusters.totalSum),
  )
}

export function computeMatrixCoefficients(points, boundingRect) {
  const { maxX, minX, maxY, minY } = boundingRect
  const xBounds = mArray(minX, maxX)
  const yBounds = mArray(minY, maxY)
  const clusters = computeClusters(points, xBounds, yBounds)
  const coefficients = clusterCoefficients(clusters, points)
  return coefficients
}
const RECT_THRESHOLD_CENTER = 0
const RECT_THRESHOLD_SIDE_VARIANCE = 0.25

function couldBeRect(points) {
  if (points.length < 4) return false

  const boundingRect = boundingCoords(points)
  const matrixCoefficients = computeMatrixCoefficients(points, boundingRect)

  let [maxC, minC] = [0, 1]
  for (let i = 0; i < MATRIX_SIZE; i++) {
    for (let j = 0; j < MATRIX_SIZE; j++) {
      if (!(i === j && j === 1)) {
        maxC = Math.max(maxC, matrixCoefficients[i][j])
        minC = Math.min(minC, matrixCoefficients[i][j])
      }
    }
  }

  if (
    matrixCoefficients[1][1] <= RECT_THRESHOLD_CENTER &&
    maxC - minC < RECT_THRESHOLD_SIDE_VARIANCE
  ) {
    return { coefficients: matrixCoefficients, boundingRect }
  }
  return undefined
}

function recognizeRect(points, rectDetectionData) {
  const { minX, minY, maxX, maxY } = rectDetectionData.boundingRect
  return {
    boundingRect: rectDetectionData.boundingRect,
    boundingPoints: [
      [minX, minY],
      [minX, maxY],
      [maxX, maxY],
      [maxX, minY],
      [minX, minY],
    ],
    shape: Shapes.rectangle,
    points,
  }
}

function recognizeLine(points) {
  if (!(points && points.length)) return {}
  const accum = Array(numAngleCells)
  const houghConfig = {
    rhoStep: points.length > 50 ? 50 : 5,
  }
  points.forEach((x) => constructHoughAccumulator(houghConfig, accum, ...x))
  const angle = findMaxInHough(accum, points.length - 1)

  if (angle !== undefined) {
    return {
      shape: Shapes.line,
      angle: 90 - angle,
      hough: accum,
      points,
    }
  }

  return {}
}

function recognizeFromPoints(points) {
  const rectDetectData = couldBeRect(points)
  if (rectDetectData) {
    return recognizeRect(points, rectDetectData)
  } else if (couldBeLine(points)) {
    return recognizeLine(points)
  }

  return {}
}

export const Shapes = {
  rectangle: "rect",
  line: "line",
}

export default recognizeFromPoints