Groovy, Guava, AWS Rekognition SDK, Indexes faces and queries indexes looking for similar faces

matchFaces

#!/usr/bin/env groovy

@Grapes([
    @Grab(group='com.amazonaws', module='aws-java-sdk-rekognition', version='1.11.117'),
    @Grab(group='com.google.guava', module='guava', version='21.0')
])

import com.amazonaws.auth.AWSCredentials
import com.amazonaws.auth.DefaultAWSCredentialsProviderChain
import com.amazonaws.auth.profile.ProfileCredentialsProvider
import com.amazonaws.regions.Regions
import com.amazonaws.services.rekognition.AmazonRekognition
import com.amazonaws.services.rekognition.AmazonRekognitionClientBuilder
import com.amazonaws.services.rekognition.model.IndexFacesRequest
import com.amazonaws.services.rekognition.model.CreateCollectionRequest
import com.amazonaws.services.rekognition.model.DeleteCollectionRequest
import com.amazonaws.services.rekognition.model.Face
import com.amazonaws.services.rekognition.model.FaceMatch
import com.amazonaws.services.rekognition.model.Image
import com.amazonaws.services.rekognition.model.ListCollectionsRequest
import com.amazonaws.services.rekognition.model.ListFacesRequest
import com.amazonaws.services.rekognition.model.SearchFacesRequest

import com.google.common.base.Stopwatch
import com.google.common.collect.HashMultimap

import java.io.FilenameFilter
import java.nio.ByteBuffer
import java.util.concurrent.TimeUnit
import java.util.List

import groovy.json.JsonOutput
import groovy.transform.Canonical

def defaultCollectionId = 'ephemeral-faces-collection'
def defaultMatchConfidenceThreshold = 80F

def printErr = System.err.&println

def cli = new CliBuilder(header: 'AWS Rekognition Match Faces POC', usage:'./matchFaces [options] <directoryOfImagesWithFaces>', width: 140, footer: 'Note: The AWS Rekognition API only supports facial recognition for JPEG and PNG image formats. This tool will examine your target directory for images with the .jpg, .jpeg, and .png file extensions and disregard the rest.')
cli.collectionId(args:1, argName:'id', "The collection id to use for this execution [${defaultCollectionId}]")
cli.delete('Delete the collection of faces post processing')
cli.forceRecreate('Force collection re-creation, if the collection exists')
cli.help('Show this menu')
cli.matchConfidenceThreshold(args:1, argName:'threshold', 'The match confidence threshold to use, values 0-100 [80]')
cli.saveImageData('Save image data as JSON for each response from AWS index image operation')
cli.verbose('Verbose output')

def cliOptions = cli.parse(args)

if (cliOptions.help || cliOptions.arguments().size() != 1) {
  cli.usage()
  System.exit(0)
}

def collectionId = cliOptions?.collectionId ?: defaultCollectionId
def matchConfidenceThreshold = cliOptions?.matchConfidenceThreshold ? Float.parseFloat(cliOptions.matchConfidenceThreshold) : defaultMatchConfidenceThreshold
def deleteCollection = cliOptions.delete
def forceRecreate = cliOptions.forceRecreate
def saveImageData = cliOptions.saveImageData
def verboseOutput = cliOptions.verbose

class ImageFilter implements FilenameFilter {
  public boolean accept(File file, String filename) {
    filename.endsWith('jpeg') || filename.endsWith('jpg') || filename.endsWith('png')
  }
}

def sourceImages = new File(cliOptions.arguments().first())

if (matchConfidenceThreshold < 0 || matchConfidenceThreshold > 100) {
  printErr "The supplied matchConfidenceThreshold of ${matchConfidenceThreshold} is invalid. Please supply a value between 0-100 or remove the argument to use the default value of 80."
  System.exit(-1)
}

if (!sourceImages.exists()) {
  printErr "The directory '${sourceImages.path}' does not exist."
  System.exit(-1)
}

def listOfImages = sourceImages.listFiles(new ImageFilter())

if (!listOfImages) {
  printErr "There are no images suitable for processing in the directory '${sourceImages.path}'"
  System.exit(0)
}

def awsRekognitionClient = AmazonRekognitionClientBuilder
   .standard()
   .withRegion(Regions.US_WEST_2)
   .withCredentials(new DefaultAWSCredentialsProviderChain())
   .build()

def collectionExists = awsRekognitionClient.listCollections(new ListCollectionsRequest()).collectionIds.contains(collectionId)

if (collectionExists && forceRecreate) {
  awsRekognitionClient.deleteCollection(new DeleteCollectionRequest().withCollectionId(collectionId))
  awsRekognitionClient.createCollection(new CreateCollectionRequest().withCollectionId(collectionId))
} else if (!collectionExists) {
  awsRekognitionClient.createCollection(new CreateCollectionRequest().withCollectionId(collectionId))
}

def existingFaces = awsRekognitionClient.listFaces(new ListFacesRequest().withCollectionId(collectionId)).faces

def existingFacesExternalImageIds = existingFaces.collect { face ->
  face.externalImageId
}

def totalProcessingStopwatch = Stopwatch.createStarted()
def perImageProcessingStopwatch = Stopwatch.createUnstarted()

listOfImages.each { file ->

  if (existingFacesExternalImageIds.contains(file.name)) {

    if (verboseOutput) {
      println "The image file ${file.name} has already been indexed. Skipping."
    }
  } else {

    perImageProcessingStopwatch.start()

    def indexFacesRequest = new IndexFacesRequest()
       .withImage(new Image()
         .withBytes(ByteBuffer.wrap(file.getBytes())))
       .withCollectionId(collectionId)
       .withExternalImageId(file.name)
       .withDetectionAttributes('ALL')

     def indexFacesResponse = awsRekognitionClient.indexFaces(indexFacesRequest)

     perImageProcessingStopwatch.stop()

     if (verboseOutput) {
       println "Found ${indexFacesResponse.faceRecords.size()} faces in image ${file.name} in ${perImageProcessingStopwatch.elapsed(TimeUnit.MILLISECONDS)} ms"
     } else if (indexFacesResponse.faceRecords.size() != 1) {
       println "Found ${indexFacesResponse.faceRecords.size()} faces in image ${file.name} in ${perImageProcessingStopwatch.elapsed(TimeUnit.MILLISECONDS)} ms"
     }

     if (saveImageData) {
       new File("${sourceImages.path}/${file.name}_results.json").write(JsonOutput.prettyPrint(JsonOutput.toJson(indexFacesResponse)))
     }

     perImageProcessingStopwatch.reset()
  }
}

totalProcessingStopwatch.stop()

if (verboseOutput) {
  println "Took a total of ${totalProcessingStopwatch.elapsed(TimeUnit.SECONDS)} seconds to index ${listOfImages.size()} images"
}

println "/----------------------------------------------------------------------------------------------------------------------------------------\\"
println "| Source Face ID                       | Source File          | Target Face ID                       | Target File          | Similarity |"
println "|----------------------------------------------------------------------------------------------------------------------------------------|"

def leftAlignedFormatOutput = "| %s | %-20s | %s | %-20s | %-10f |%n"

existingFaces.each { face ->

    awsRekognitionClient
      .searchFaces(new SearchFacesRequest()
        .withCollectionId(collectionId)
        .withFaceId(face.faceId)
        .withFaceMatchThreshold(matchConfidenceThreshold))
      .faceMatches.each { faceMatch ->

        System.out.format(leftAlignedFormatOutput, face.faceId, face.externalImageId, faceMatch.face.faceId, faceMatch.face.externalImageId, faceMatch.similarity)
      }
  }

println "\\----------------------------------------------------------------------------------------------------------------------------------------/"

if (deleteCollection) {
  awsRekognitionClient.deleteCollection(new DeleteCollectionRequest().withCollectionId(collectionId))
}