app.py

"""
Jialu Bi
CS5330 Lab 1
Generator for ASCII art
"""

import cv2
import numpy as np
import matplotlib.pyplot as plt
import gradio as gr
from skimage.metrics import structural_similarity as ssim
from skimage.metrics import mean_squared_error

# A class for ASCII genrator related functions
class ASCIIGenerator:
    def __init__(self):
        self.ASCIIChars = ['#', '@', '%', '*', '+', '-', '.', ' ']
        # self.ASCIIChars = ['█', '▓', '▒', '░', ' ']

    # Resize the image to make sure consistent result
    def resizeImg(self, img):
        h, w, c = img.shape
        aspectRatio = w / h
        imgResized = cv2.resize(img, (int(160 * aspectRatio), 160))
        return imgResized

    # Convert to gray scale
    # The input should be the resized image
    def cvtGrayScale(self, imgResized):
        R, G, B = cv2.split(imgResized)
        R, G, B = R.astype(np.float32), G.astype(np.float32), B.astype(np.float32)

        # Apply weights for more refined control
        imgGray = 0.299 * R + 0.587 * G + 0.114 * B
        imgGray = np.clip(imgGray, 0, 255).astype(np.uint8)
        return imgGray

    # Preprocess the image
    # The input should be the gray scale image
    def preprocess(self, imgGray):
        # Histogram equalization
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        imgEqualized = clahe.apply(imgGray)

        # Increase contrast
        imgEqualized = cv2.convertScaleAbs(imgEqualized, alpha=1.6, beta=0)

        # Detect the edges and emphasize them
        edges = cv2.Canny(imgEqualized, threshold1=50, threshold2=150)
        invertedEdges = cv2.bitwise_not(edges)
        imgEnhanced = cv2.addWeighted(imgGray, 0.7, invertedEdges, 0.3, 0)
        return imgEnhanced

    # Map to ASCII Characters
    # The input should be the preprocessed image
    def mapASCII(self, imgEnhanced):
        ASCIIArt = ''
        h, w = imgEnhanced.shape

        # Iterate through the pixels to create the ASCII art
        for i in range(h):
            for j in range(w):
                pixelVal = imgEnhanced[i, j]
                # Assign a proper ASCII character according to the pixel's intensity value
                ASCIIArt += self.ASCIIChars[pixelVal * len(self.ASCIIChars) // 256]
            ASCIIArt += '\n'
        return ASCIIArt

    # Main function for generating ASCII art
    def generateASCIIArt(self, img):
        imgResized = self.resizeImg(img)
        imgGray = self.cvtGrayScale(imgResized)
        imgEnhanced = self.preprocess(imgGray)
        ASCIIArt = self.mapASCII(imgEnhanced)
        return ASCIIArt

    # Convert ASCII art back to grayscale for performance evaluation
    def ASCIItoGray(self, ASCIIArt):
        # Create a dictionary to store the corresponding intensity of ASCII characters
        ASCIICharsVal = {}
        offset = 0
        for char in self.ASCIIChars:
            ASCIICharsVal[char] = offset
            offset += 256 // len(self.ASCIIChars)

        # Map ASCII characters back to pixels
        ASCIIGray = []
        lines = ASCIIArt.split('\n')[:-1]
        for line in lines:
            row = []
            for char in line:
                row.append(ASCIICharsVal[char])
            ASCIIGray.append(row)
        ASCIIGray = np.array(ASCIIGray, dtype=np.uint8)

        return ASCIIGray

    # Evaluate the similarity between the grayscale version of the original image and its ASCII art version
    def compare(self, ASCIIGray, imgGray):
        # Compute SSIM
        ssimVal, _ = ssim(imgGray, ASCIIGray, full=True)
        # Compute MSE
        mseVal = mean_squared_error(imgGray, ASCIIGray)
        # Compute PSNR
        psnrVal = cv2.PSNR(imgGray, ASCIIGray)
        return ssimVal, mseVal, psnrVal


# Gradio interface function
def generateArt(img):
    if img is None:
        return None, None, None

    # Generate the ASCII art
    generator = ASCIIGenerator()
    ASCIIArt = generator.generateASCIIArt(img)

    # Evaluate the performance of the algorithm
    imgResized = generator.resizeImg(img)
    imgGray = generator.cvtGrayScale(imgResized)
    ASCIIGray = generator.ASCIItoGray(ASCIIArt)
    scores = generator.compare(ASCIIGray, imgGray)
    report = f'SSIM: {round(scores[0], 2)}, MSE: {round(scores[1], 2)}, PSNR: {round(scores[2], 2)}'

    # Save the ASCII art to a text file
    outputPath = 'ascii_art.txt'
    with open(outputPath, 'w', encoding='utf-8') as file:
        file.write(ASCIIArt)
    return f'<pre style="font-size: 8px; line-height: 0.62;">{ASCIIArt}</pre>', outputPath, report

# Create a Gradio interface
with gr.Blocks() as demo:
    gr.Markdown('## ASCII Art Generator\nUpload an image to generate ASCII art.')
    imgInput = gr.Image(type='numpy', label='Upload Image')
    ASCIIOutput = gr.HTML(label='ASCII Art Output')
    fileOutput = gr.File(label='Download ASCII Art')
    scoreOutput = gr.Textbox(label='ASCII Conversion Performance')
    imgInput.change(fn=generateArt, inputs=imgInput, outputs=[ASCIIOutput, fileOutput, scoreOutput])


demo.launch(share=True)