ascii_image_transformer.py

#!/usr/bin/env python
'''
Convert an image to ASCII art
'''
import argparse
from PIL import Image

# pylint: disable=redefined-outer-name
def main(args):
    '''
    Main function
    '''
    image = Image.open(args.image)

    ascii_art = trim_whitespace(
        convert_image_to_ascii(image, args.width)
    )

    write_to_file(args.image, ascii_art)

    print(ascii_art)
    return ascii_art


def write_to_file(image_path, ascii_art):
    '''
    Write the ASCII art to a text file
    '''
    ascii_art_path = f'{image_path.split(".")[0]}.txt'

    with open(ascii_art_path, mode='w', encoding='utf-8') as f:
        f.write(ascii_art)


def rescale_image(image, new_width=100):
    '''
    Resizes an image preserving the aspect ratio.
    '''
    (original_width, original_height) = image.size

    aspect_ratio = original_height/float(original_width)
    new_height   = int(aspect_ratio * new_width)

    resized_image = image.resize((new_width, new_height))

    return resized_image


def grayscale_image(image):
    '''
    Convert an image to grayscale
    '''
    return image.convert("L")


def map_pixels_to_ascii_chars(image, range_width=25):
    '''
    Maps each pixel to an ASCII char based on the range in which it lies.
    '''
    ascii_chars = ["@", "#", "S", "Z", "?", "*", "=", ";", ":", ".", " "]

    pixels_in_image = list(image.getdata())

    pixels_to_chars = [
        ascii_chars[pixel_value//range_width] for pixel_value in pixels_in_image
    ]

    return "".join(pixels_to_chars)


def convert_image_to_ascii(image, new_width=50):
    '''
    Convert the image to ASCII characters
    '''
    image = grayscale_image(
        rescale_image(image, new_width)
    )

    pixels_to_chars = map_pixels_to_ascii_chars(image)
    len_pixels_to_chars = len(pixels_to_chars)

    image_ascii = [
        # pylint: disable=line-too-long
        pixels_to_chars[index: index + new_width] for index in range(0, len_pixels_to_chars, new_width)
    ]

    return "\n".join(image_ascii)


def trim_whitespace(ascii_art):
    '''
    Remove trailing whitespace from each line
    '''
    lines = ascii_art.split('\n')
    lines = [line.rstrip() for line in lines]

    return '\n'.join([
        line for line in lines if line
    ])


def parse_args():
    '''
    Parse command-line arguments
    '''
    parser = argparse.ArgumentParser(description='Convert image to ASCII art')

    parser.add_argument(
        '-i', '--image',
        type=str,
        required=True,
        help='Path to the image file'
    )

    parser.add_argument(
        '-w', '--width',
        type=int,
        default=50,
        help='Width of the ASCII art'
    )

    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()
    main(args=args)