img2img_color.py

"""
@original_author: Viet Nguyen <nhviet1009@gmail.com>
@modified_by: Daniel Song <liucan01234@gmail.com> 
Repository: https://github.com/Liu8Can/ASCII-generator
This project is open-source and protected under the GPL-3.0 License. 
Please keep this header intact and share improvements.
"""

import argparse

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageOps
from utils import get_data


def get_args():
    parser = argparse.ArgumentParser("Image to ASCII")
    parser.add_argument(
        "--input", type=str, default="data/input.jpg", help="Path to input image"
    )
    parser.add_argument(
        "--output", type=str, default="data/output.jpg", help="Path to output text file"
    )
    parser.add_argument("--language", type=str, default="english")
    parser.add_argument("--mode", type=str, default="standard")
    parser.add_argument(
        "--background",
        type=str,
        default="black",
        choices=["black", "white"],
        help="background's color",
    )
    parser.add_argument(
        "--num_cols",
        type=int,
        default=300,
        help="number of character for output's width",
    )
    parser.add_argument("--scale", type=int, default=2, help="upsize output")
    args = parser.parse_args()
    return args


def main(opt):
    if opt.background == "white":
        bg_code = (255, 255, 255)
    else:
        bg_code = (0, 0, 0)
    char_list, font, sample_character, scale = get_data(opt.language, opt.mode)
    num_chars = len(char_list)
    num_cols = opt.num_cols
    image = cv2.imread(opt.input, cv2.IMREAD_COLOR)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    height, width, _ = image.shape
    cell_width = width / opt.num_cols
    cell_height = scale * cell_width
    num_rows = int(height / cell_height)
    if num_cols > width or num_rows > height:
        print("Too many columns or rows. Use default setting")
        cell_width = 6
        cell_height = 12
        num_cols = int(width / cell_width)
        num_rows = int(height / cell_height)
    char_width, char_height = (
        font.getbbox(sample_character)[2],
        font.getbbox(sample_character)[3],
    )
    out_width = char_width * num_cols
    out_height = scale * char_height * num_rows
    out_image = Image.new("RGB", (out_width, out_height), bg_code)
    draw = ImageDraw.Draw(out_image)
    for i in range(num_rows):
        for j in range(num_cols):
            partial_image = image[
                int(i * cell_height) : min(int((i + 1) * cell_height), height),
                int(j * cell_width) : min(int((j + 1) * cell_width), width),
                :,
            ]
            partial_avg_color = np.sum(np.sum(partial_image, axis=0), axis=0) / (
                cell_height * cell_width
            )
            partial_avg_color = tuple(partial_avg_color.astype(np.int32).tolist())
            char = char_list[
                min(int(np.mean(partial_image) * num_chars / 255), num_chars - 1)
            ]
            draw.text(
                (j * char_width, i * char_height),
                char,
                fill=partial_avg_color,
                font=font,
            )

    if opt.background == "white":
        cropped_image = ImageOps.invert(out_image).getbbox()
    else:
        cropped_image = out_image.getbbox()
    out_image = out_image.crop(cropped_image)
    out_image.save(opt.output)


if __name__ == "__main__":
    opt = get_args()
    main(opt)