app.py

"""
CS5330 Fall 2024
Author: Calvin Lo
Lab 1: ASCII Art
This program generates the ASCII art representation of
an input image and compares the result to the original image
using SSIM.
"""

import gradio as gr
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageFilter, ImageEnhance
from skimage.metrics import structural_similarity as ssim

# ASCII characters to map pixel intensity
ASCII_CHARS = [
    ' ', '.', ',', ':', ';', 'i', 'l', '!', 'I', '1', 't', 'o', 'r', 'x', 'z', 'v', 
    'u', 'n', 'm', 'w', 'Q', 'B', 'N', 'M', '@'
]

def preprocess_image(image, new_width=150):
    """
    This function performs the following operations on the input image:
    -enhance contrast
    -perform edge detection (Sobel)
    -denoise edges
    -enhance contrast
    -resize image
    """
    # Open the image and convert to grayscale
    image = Image.fromarray(image)
    image = image.convert("L")

    # Contrast limited adaptive histogram equalization
    opencv_image = np.array(image)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    equalized_image = clahe.apply(opencv_image)

    # Apply Sobel edge detection
    sobelx = cv2.Sobel(equalized_image, cv2.CV_64F, 1, 0, ksize=5)
    sobely = cv2.Sobel(equalized_image, cv2.CV_64F, 0, 1, ksize=5)
    edges = np.hypot(sobelx, sobely)  # Combine gradients

    # Normalize the result to fit into the 0-255 grayscale range
    edges = np.uint8(255 * edges / np.max(edges))

    # Improve contrast of edges
    edges_image = Image.fromarray(edges)
    enhancer = ImageEnhance.Contrast(edges_image)
    enhanced_edges = enhancer.enhance(1.5)
    enhanced_edges = np.array(enhanced_edges)
    
    # Resize the image for ASCII aspect ratio
    width, height = image.size
    aspect_ratio = height / width * 0.5  # Adjusting for ASCII aspect ratio
    new_height = int(aspect_ratio * new_width)
    resized_image = Image.fromarray(enhanced_edges).resize((new_width, new_height))

    return resized_image


def image_to_ascii(image):
    """
    This function generates an ASCII representation of images.
    """
    pixels = np.array(image)
    ascii_str = ""
    for row in pixels:
        for pixel in row:
            # Mapping pixel intensity to ASCII
            ascii_str += ASCII_CHARS[pixel // 10]
        ascii_str += "\n"
    return ascii_str


def ascii_to_image(ascii_art, font_size=12, output_file='ascii_art.png'):
    """"
    This function converts ASCII text art to a png image.
    """

    # Specify font
    font_path = "DejaVuSansMono.ttf"
    font = ImageFont.truetype(font_path, font_size)

    # Get the required image dimensions
    lines = ascii_art.splitlines()
    max_width = max(font.getbbox(line)[2] for line in lines)  # Get max width
    img_height = len(lines) * font_size

    # Generate image with white background
    image = Image.new("RGB", (max_width, img_height), "white")
    draw = ImageDraw.Draw(image)

    # Draw each character from the ASCII text on the image
    for y, line in enumerate(lines):
        draw.text((0, y * font_size), line, fill="black", font=font)

    image.save(output_file)
    return image


def compare_ssim(image1, image2):
    """
    This function generates the structural similarity index
    measure for the two input images.
    """
    # Convert original color image to grayscale
    image1 = cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY)
    
    # Convert ASCII image to grayscale and cv2 format
    image2 = np.array(image2.convert("L"))

    # Resize so both images have the same dimensions.
    image2 = cv2.resize(image2, (image1.shape[1], image1.shape[0]))

    score, _ = ssim(image1, image2, full=True)
    return score


def process_image(image):
    """
    This function takes an image and generates:
        -ASCII art as a string
        -ASCII text file
        -Compute SSIM for the ASCII art image and the original image
    """
    # Get image edges and convert to ASCII
    edges = preprocess_image(image)
    ascii_art = image_to_ascii(edges)

    # Convert ASCII art back to an image for SSIM calculation
    ascii_image = ascii_to_image(ascii_art)

    # Calculate SSIM
    ssim_score = compare_ssim(image, ascii_image)

    # Write ASCII art to txt file
    ascii_txt_path = "ascii_art_output.txt"
    with open(ascii_txt_path, "w") as f:
        f.write(ascii_art)

    return ascii_art, ssim_score, ascii_txt_path


def gradio_interface(image):
    """
    This function generates the interface outputs for a given input image.
    The outputs include the formatted SSIM, ASCII art, and ASCII art txt file path.
    """
    ascii_art, ssim_score, ascii_txt_path = process_image(image)
    ssim_formatted = f'<div style="font-size: 20px;">The structural similarity index measure is {ssim_score:.3f}</div>'
    # Wrap the ASCII art in HTML with custom font size
    styled_ascii_art = f'<pre style="font-family: monospace; font-size: 11px;">{ascii_art}</pre>'
    
    return styled_ascii_art, ssim_formatted, ascii_txt_path


def main():
    # Setup Gradio interface
    with gr.Blocks() as interface:
        # Add title and description
        gr.Markdown("<h1 style='font-size: 40px;'>Image to ASCII Art Converter</h1>")
        gr.Markdown("<p style='font-size: 18px;'>Upload an image, and this tool will generate an ASCII art version of the image. The Structural Similarity Index (SSIM) is calculated to evaluate the quality as a performance metric.</p>")

        # Image input and ASCII art + SSIM score output
        image_input = gr.Image(label="Upload an image")
        ascii_art_output = gr.HTML()
        ssim_score_output = gr.HTML()

        # File output for downloading, initially hidden
        ascii_file_output = gr.File(label="Download ASCII Art")

        # Update interface when image is uploaded
        image_input.change(fn=gradio_interface,
                           inputs=image_input,
                           outputs=[ascii_art_output, ssim_score_output, ascii_file_output])

    # Launch interface
    interface.launch()


if __name__ == "__main__":
    main()