<commit_before>import * as nj            from 'numjs'

import {
  FaceEmbedding,
  log,
  VERSION,
}                         from './config'
import {
  Face,
}                         from './face'
import {
  distance,
  imageToData,
  loadImage,
  resizeImage,
}                         from './misc'

import {
  PythonFacenet,
}                         from './python3/python-facenet'

// Interface for Cache
export interface Alignable {
  align(imageData: ImageData | string): Promise<Face[]>,
}

// Interface for Cache
export interface Embeddingable {
  embedding(face: Face): Promise<FaceEmbedding>,
}

export class Facenet implements Alignable, Embeddingable {
  private pythonFacenet: PythonFacenet

  constructor() {
    log.verbose('Facenet', `constructor() v${VERSION}`)
    this.pythonFacenet = new PythonFacenet()
  }

  public async init(): Promise<void> {
    await this.initFacenet()
    await this.initMtcnn()
  }

  public async initFacenet(): Promise<void> {
    log.verbose('Facenet', 'initFacenet()')
    const start = Date.now()
    await this.pythonFacenet.initFacenet()
    log.verbose('Facenet', 'initFacenet() cost %d milliseconds', Date.now() - start)
  }

  public async initMtcnn(): Promise<void> {
    log.verbose('Facenet', 'initMtcnn()')
    const start = Date.now()
    await this.pythonFacenet.initMtcnn()
    log.verbose('Facenet', 'initMtcnn() cost %d milliseconds', Date.now() - start)
  }

  public async quit(): Promise<void> {
    log.verbose('Facenet', 'quit()')
    await this.pythonFacenet.quit()
  }

  /**
   * Alignment the image, get faces list
   * @param image
   */
  public async align(imageData: ImageData | string): Promise<Face[]> {
    if (typeof imageData === 'string') {
      log.verbose('Facenet', 'align(%s)', imageData)
      const image = await loadImage(imageData)
      imageData = imageToData(image)
    } else {
      log.verbose('Facenet', 'align(%dx%d)', imageData.width, imageData.height)
    }

    const [boundingBoxes, landmarks] = await this.pythonFacenet.align(imageData)
    log.silly('Facenet', 'align() pythonFacenet.align() done: %s', boundingBoxes)

    const xyLandmarks = this.transformMtcnnLandmarks(landmarks)

    const faceList: Face[] = []
    for (const i in boundingBoxes) {
      const boundingBox = this.squareBox(boundingBoxes[i])
      const confidence = boundingBoxes[i][4]
      const marks = xyLandmarks[i]

      // boundary out of image
      if (boundingBox[0] < 0 || boundingBox[1] < 0
        ||  boundingBox[2] > imageData.width
        ||  boundingBox[3] > imageData.height
      ) {
        log.silly('Facenet', 'align(%dx%d) box[%s] out of boundary, skipped',
                              imageData.width, imageData.height,
                              boundingBox)
        continue
      }

      const face = new Face(imageData, boundingBox)
      face.init(marks, confidence)
      // face.confidence(confidence)
      // face.landmarks(marks)
      faceList.push(face)
    }

    return faceList
  }

  /**
   * Get the 128 dims embeding from image(s)
   */
  public async embedding(face: Face): Promise<FaceEmbedding> {
    log.verbose('Facenet', 'embedding(%s)', face)

    let imageData = face.imageData
    if (imageData.width !== imageData.height) {
      log.warn('Facenet', 'embedding(%s) %dx%d not square!',
                          face, imageData.width, imageData.height)
      throw new Error('should be a square image because it will be resized to 160x160')
    }

    if (imageData.width !== 160) {
      log.verbose('Facenet', 'embedding(%dx%d) got a face not 160x160, resizing...',
                            imageData.width, imageData.height)
      imageData = await resizeImage(imageData, 160, 160)
    }

    const embedding = await this.pythonFacenet.embedding(imageData)

    // Set embedding to face
    face.embedding =  nj.array(embedding)

    return face.embedding
  }

  public transformMtcnnLandmarks(landmarks: number[][]): number[][][] {
    // landmarks has a strange data structure:
    // https://github.com/kpzhang93/MTCNN_face_detection_alignment/blob/bace6de9fab6ddf41f1bdf1c2207c50f7039c877/code/codes/camera_demo/test.m#L70
    const tLandmarks = nj.array(landmarks.reduce((a, b) => a.concat(b), []))
                          .reshape(10, -1)
                          .T as nj.NdArray<number>

    const faceNum = tLandmarks.shape[0]

    const xyLandmarks = nj.zeros(tLandmarks.shape)
    const xLandmarks = xyLandmarks.slice(null as any, [null,  xyLandmarks.shape[1], 2] as any)
    const yLandmarks = xyLandmarks.slice(null as any, [1,     xyLandmarks.shape[1], 2] as any)
    xLandmarks.assign(
      tLandmarks.slice(null as any, [null, 5] as any),
      false,
    )
    yLandmarks.assign(
      tLandmarks.slice(null as any, [5, 10] as any),
      false,
    )
<commit_msg>BREAKING CHANGES: facenet.embedding(face) will not modify face directly<commit_after>import * as nj            from 'numjs'

import {
  FaceEmbedding,
  log,
  VERSION,
}                         from './config'
import {
  Face,
}                         from './face'
import {
  distance,
  imageToData,
  loadImage,
  resizeImage,
}                         from './misc'

import {
  PythonFacenet,
}                         from './python3/python-facenet'

// Interface for Cache
export interface Alignable {
  align(imageData: ImageData | string): Promise<Face[]>,
}

// Interface for Cache
export interface Embeddingable {
  embedding(face: Face): Promise<FaceEmbedding>,
}

export class Facenet implements Alignable, Embeddingable {
  private pythonFacenet: PythonFacenet

  constructor() {
    log.verbose('Facenet', `constructor() v${VERSION}`)
    this.pythonFacenet = new PythonFacenet()
  }

  public async init(): Promise<void> {
    await this.initFacenet()
    await this.initMtcnn()
  }

  public async initFacenet(): Promise<void> {
    log.verbose('Facenet', 'initFacenet()')
    const start = Date.now()
    await this.pythonFacenet.initFacenet()
    log.verbose('Facenet', 'initFacenet() cost %d milliseconds', Date.now() - start)
  }

  public async initMtcnn(): Promise<void> {
    log.verbose('Facenet', 'initMtcnn()')
    const start = Date.now()
    await this.pythonFacenet.initMtcnn()
    log.verbose('Facenet', 'initMtcnn() cost %d milliseconds', Date.now() - start)
  }

  public async quit(): Promise<void> {
    log.verbose('Facenet', 'quit()')
    await this.pythonFacenet.quit()
  }

  /**
   * Alignment the image, get faces list
   * @param image
   */
  public async align(imageData: ImageData | string): Promise<Face[]> {
    if (typeof imageData === 'string') {
      log.verbose('Facenet', 'align(%s)', imageData)
      const image = await loadImage(imageData)
      imageData = imageToData(image)
    } else {
      log.verbose('Facenet', 'align(%dx%d)', imageData.width, imageData.height)
    }

    const [boundingBoxes, landmarks] = await this.pythonFacenet.align(imageData)
    log.silly('Facenet', 'align() pythonFacenet.align() done: %s', boundingBoxes)

    const xyLandmarks = this.transformMtcnnLandmarks(landmarks)

    const faceList: Face[] = []
    for (const i in boundingBoxes) {
      const boundingBox = this.squareBox(boundingBoxes[i])
      const confidence = boundingBoxes[i][4]
      const marks = xyLandmarks[i]

      // boundary out of image
      if (boundingBox[0] < 0 || boundingBox[1] < 0
        ||  boundingBox[2] > imageData.width
        ||  boundingBox[3] > imageData.height
      ) {
        log.silly('Facenet', 'align(%dx%d) box[%s] out of boundary, skipped',
                              imageData.width, imageData.height,
                              boundingBox)
        continue
      }

      const face = new Face(imageData, boundingBox)
      face.init(marks, confidence)
      // face.confidence(confidence)
      // face.landmarks(marks)
      faceList.push(face)
    }

    return faceList
  }

  /**
   * Get the 128 dims embeding from image(s)
   */
  public async embedding(face: Face): Promise<FaceEmbedding> {
    log.verbose('Facenet', 'embedding(%s)', face)

    let imageData = face.imageData
    if (imageData.width !== imageData.height) {
      log.warn('Facenet', 'embedding(%s) %dx%d not square!',
                          face, imageData.width, imageData.height)
      throw new Error('should be a square image because it will be resized to 160x160')
    }

    if (imageData.width !== 160) {
      log.verbose('Facenet', 'embedding(%dx%d) got a face not 160x160, resizing...',
                            imageData.width, imageData.height)
      imageData = await resizeImage(imageData, 160, 160)
    }

    const embedding = await this.pythonFacenet.embedding(imageData)

    // Set embedding to face
    // face.embedding =  nj.array(embedding)

    return nj.array(embedding)
  }

  public transformMtcnnLandmarks(landmarks: number[][]): number[][][] {
    // landmarks has a strange data structure:
    // https://github.com/kpzhang93/MTCNN_face_detection_alignment/blob/bace6de9fab6ddf41f1bdf1c2207c50f7039c877/code/codes/camera_demo/test.m#L70
    const tLandmarks = nj.array(landmarks.reduce((a, b) => a.concat(b), []))
                          .reshape(10, -1)
                          .T as nj.NdArray<number>

    const faceNum = tLandmarks.shape[0]

    const xyLandmarks = nj.zeros(tLandmarks.shape)
    const xLandmarks = xyLandmarks.slice(null as any, [null,  xyLandmarks.shape[1], 2] as any)
    const yLandmarks = xyLandmarks.slice(null as any, [1,     xyLandmarks.shape[1], 2] as any)
    xLandmarks.assign(
      tLandmarks.slice(null as any, [null, 5] as any),
      false,
    )
    yLandmarks.assign(
      tLandmarks.slice(null as any, [5, 10] as any),
      false,
    )

{'fraction_non_alphanumeric': 0.12769895126465144, 'fraction_numerical': 0.01418877236273905, 'mean_word_length': 2.7182721712538225, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 27, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4575439', 'n_tokens_mistral': 3286, 'n_tokens_neox': 3115, 'n_words': 707}