Playing with AWS Rekognition

rekognition.py

"""Playing with AWS Rekognition."""
import argparse
import hashlib
import json
import logging

import boto3


logging.basicConfig(
    level=logging.INFO,
    format="[%(asctime)s]:%(name)s:%(message)s"
)


class RekognitionExample:
    """Create a client and a method to get labels for a given image."""
    @staticmethod
    def get_file_name(**kwargs):
        """Get a file name unique to this request."""
        contents = json.dumps(kwargs).encode()
        digest = hashlib.md5(contents).hexdigest()

        return f"{digest}.json"

    def __init__(self, client=None, logger=None, **kwargs):
        """Override."""
        self.client = client if client \
            else boto3.client("rekognition")
        self.logger = logger if logger \
            else logging.getLogger(self.__class__.__name__)

        self.get_labels(**kwargs)

    def get_labels(self, **kwargs):
        """Get a list of labels for an image."""
        if not kwargs:
            self.labels = list()
            return

        file_name = self.get_file_name(**kwargs)
        response = dict()

        self.logger.info(f"GetLabels:{kwargs}")

        try:
            with open(file_name) as in_file:
                response = json.load(in_file)

            self.logger.info(f"Read:{file_name}")
        except FileNotFoundError:
            response = self.client.detect_labels(**kwargs)

            with open(file_name, "w") as out_file:
                json.dump(response, out_file)

            self.logger.info(f"Write:{file_name}")

        self.labels = [label["Name"] for label in response["Labels"]]
        self.logger.info(f"Labels:{self.labels}")


def parse_args():
    """Parse input arguments."""
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "Bucket", help="S3 Bucket"
    )
    parser.add_argument(
        "Name", help="S3 Key name"
    )
    parser.add_argument(
        "-c", "--min-confidence",
        dest="MinConfidence",
        type=int,
        default=55,
        help="Minimum confidence"
    )

    return vars(parser.parse_args())


def main():
    arguments = parse_args()
    rekognition = RekognitionExample(
        Image={
            "S3Object": {
                "Bucket": arguments.pop("Bucket"),
                "Name": arguments.pop("Name")
            }
        },
        **arguments
    )

    print(rekognition.labels)


if __name__ == "__main__":
    main()