ASCII_functions.py

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

# Constants
DEFAULT_WIDTH = 150
CONTRAST_FACTOR = 2.0
ASCII_CHARS = "@B%8&WM#*oahkbdpqwmZO0QLCJUYXzcvunxrjft/|)(1}{][?-_+~i!lI;:,\"^`'. "
ASCII_CHARS = list(ASCII_CHARS)


def resize_image(image, new_width=DEFAULT_WIDTH, resampling_filter=Image.LANCZOS):
    """Resize the image while maintaining the aspect ratio."""
    width, height = image.size
    aspect_ratio = height / width
    new_height = int(new_width * aspect_ratio)
    resized_image = image.resize((new_width, new_height), resampling_filter)

    # Apply a sharpening filter after resizing
    sharpened_image = resized_image.filter(ImageFilter.SHARPEN)

    return sharpened_image


def convert_to_grayscale(image):
    """Convert the image to grayscale using weighted sum of RGB channels."""
    # Get the image data
    pixels = image.load()
    width, height = image.size

    # Check if the image is already in grayscale mode
    if image.mode == 'L':
        return image

    # Create a new grayscale image
    grayscale_image = Image.new("L", (width, height))
    grayscale_pixels = grayscale_image.load()

    # Define the weights for each channel
    weights = (0.299, 0.587, 0.114)

    # Process each pixel
    for x in range(width):
        for y in range(height):
            r, g, b = pixels[x, y][:3]  # Get RGB values
            # Calculate the grayscale value using the weighted sum
            gray = int(weights[0] * r + weights[1] * g + weights[2] * b)
            grayscale_pixels[x, y] = gray

    return grayscale_image

def adjust_contrast(image, factor=CONTRAST_FACTOR):
    """Adjust the contrast of the image."""
    enhancer = ImageEnhance.Contrast(image)
    return enhancer.enhance(factor)


def apply_edge_detection(image):
    """Apply edge detection to the image using Canny algorithm with preprocessing and automatic threshold detection."""

    # Convert image to numpy array
    image_np = np.array(image)

    # Print an error message if the image is not in grayscale
    if len(image_np.shape) == 3:
        raise ValueError("Edge detection requires a grayscale image.")

    # Apply Gaussian blur to reduce noise
    blurred_image = cv2.GaussianBlur(image_np, (5, 5), 1.4)

    # Compute the median of the pixel intensities
    median_intensity = np.median(blurred_image)

    # Set the lower and upper thresholds for Canny edge detection
    lower_threshold = int(max(0, (1.0 - 0.2) * median_intensity))
    upper_threshold = int(min(255, (1.0 + 0.2) * median_intensity))

    # Apply Canny edge detection with the computed thresholds
    edges = cv2.Canny(blurred_image, threshold1=lower_threshold, threshold2=upper_threshold)

    # Convert the result back to PIL Image and return
    edge_image = Image.fromarray(edges)

    return edge_image


def map_pixels_to_ascii(image):
    """Map the pixel values of the image to ASCII characters."""
    pixels = image.getdata()
    ascii_str = ""
    num_chars = len(ASCII_CHARS)

    for pixel in pixels:
        index = pixel * (num_chars - 1) // 255
        ascii_str += ASCII_CHARS[index]

    return ascii_str

def format_ascii_art(ascii_str, width):
    """Format the ASCII string into a structured ASCII art."""
    ascii_art = ""
    for i in range(0, len(ascii_str), width):
        ascii_art += ascii_str[i:i+width] + "\n"
    return ascii_art

def process_image(image):
    """Process the input image to prepare it for ASCII art generation."""
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    grayscale_image = convert_to_grayscale(image)
    adjusted_image = adjust_contrast(grayscale_image)
    adjusted_image = resize_image(adjusted_image)
    return adjusted_image

def generate_ascii_art(image):   
    """Generate ASCII art from the given image."""
    # Ensure the input is a PIL Image object
    adjusted_image = process_image(image)
    ascii_str = map_pixels_to_ascii(adjusted_image)
    ascii_art = format_ascii_art(ascii_str, adjusted_image.width)
    return ascii_art

def ascii_to_image(ascii_art, width):
    """Convert ASCII art back to a grayscale image."""
    ascii_chars = ASCII_CHARS
    num_chars = len(ascii_chars)

    # Split the ASCII art into lines
    ascii_lines = ascii_art.split('\n')
    height = len(ascii_lines)

    # Create a blank image
    image = np.zeros((height, width), dtype=np.uint8)

    for y, line in enumerate(ascii_lines):
        for x, char in enumerate(line):
            if char in ascii_chars:
                index = ascii_chars.index(char)
                pixel_value = index * 255 // (num_chars - 1)
                image[y, x] = pixel_value

    # Convert the numpy array to a PIL Image
    image = Image.fromarray(image)
    image = adjust_contrast(image)
    return image

def calculate_similarity(original_image, ascii_art):
   # Convert ASCII art back to image
    ascii_image = ascii_to_image(ascii_art, DEFAULT_WIDTH)

    # Ensure ascii_image is a PIL Image object
    if not isinstance(ascii_image, Image.Image):
        ascii_image = Image.fromarray(ascii_image)

    # convert original image to grayscale
    original_image = convert_to_grayscale(original_image)

    # Resize the original image to the desired width
    original_resize = original_image.resize((ascii_image.width, ascii_image.height), Image.LANCZOS)
    original_resize = original_resize.filter(ImageFilter.SHARPEN)

    # Apply edge detection
    original_image_edge = apply_edge_detection(original_resize)
    # convert the original image to binary
    threshold = 0
    original_image_edge = original_image_edge.point(lambda p: p > threshold and 255)

    ascii_image_edge = apply_edge_detection(ascii_image)
    ascii_image_edge = ascii_image_edge.point(lambda p: p > threshold and 255)

    # Convert images to numpy arrays
    original_array = np.array(original_image_edge)
    ascii_array = np.array(ascii_image_edge)

    # Check if the dimensions match
    if original_array.shape != ascii_array.shape:
        raise ValueError(f"Dimension mismatch: Original Image {original_array.shape}, ASCII Image {ascii_array.shape}")

    # Calculate SSIM
    ssim_value = ssim(original_array, ascii_array)

    # Calculate MSE
    mse = np.mean((original_array - ascii_array) ** 2)


    return ssim_value, mse