desaturation_computation.py

import sys
from PIL import Image
import json
from util import DICT_NAME, round_base, hash_color, savergb
import numpy as np
from time import time

with open(DICT_NAME, "r") as file:
    deactivation_dictionary = json.load(file)

# default is to 2k
def scale_image(input_filename, output_filename="scaled_image.png", width=2048, height=1080):
    # Open the input image
    img = Image.open(input_filename).convert('RGB')

    # Calculate new size to maintain aspect ratio
    img_ratio = img.width / img.height
    target_ratio = width / height
    if img_ratio > target_ratio:
        # Width is the limiting factor
        new_width = width
        new_height = int(width / img_ratio)
    else:
        # Height is the limiting factor
        new_height = height
        new_width = int(height * img_ratio)

    # Resize the image
    img_resized = img.resize((new_width, new_height), Image.LANCZOS)

    # Create a new white image to place the resized image on
    final_img = Image.new('RGB', (width, height), 'white')
    final_img.paste(img_resized, ((width - new_width) // 2, (height - new_height) // 2))

    # Convert RGB to CMY (for storage, not output)
    # cmy_values = [(255 - pixel[0], 255 - pixel[1], 255 - pixel[2]) for pixel in final_img.getdata()]

    # save the scaled image for debug purpose
    final_img.save(output_filename)

    return final_img.getdata()

# input: target CMY value
def compute_desaturation(rgb_values, tarr_save_name='out.tarr', img_save_name='out.png', gen_preview=True, yield_progress=True):
    width, height = rgb_values.size
    tarr = np.empty((4, height, width), dtype=np.int16)
    if gen_preview:
        newimg = np.empty((height, width, 3))
        
    for i in range(height):
        if yield_progress:
            if i % max(1, height // 30) == 0:
                yield(i / height * 0.8)
        
        for j in range(width):
            pixel = rgb_values.getpixel((j, i))
            times, real_rgb = deactivation_dictionary[hash_color(*map(round_base, pixel))]
            if gen_preview:
                newimg[i, j] = real_rgb
            
            tarr[0, i, j] = times[0]
            tarr[1, i, j] = times[1]
            tarr[2, i, j] = times[2]
            tarr[3, i, j] = times[3]
                  
    if yield_progress:      
        yield(0.8)
            
    with open(tarr_save_name, 'w+') as fout:
        json.dump(tarr.tolist(), fout)
        print("%s created" % tarr_save_name) 
    
    if yield_progress:  
        yield(0.95)
            
    if gen_preview:
        savergb(newimg, img_save_name)
        print("%s saved" % img_save_name)
    
    if yield_progress:  
        yield(1)
    
def run_desaturation_computation(filename):
    for progress in compute_desaturation(Image.open(filename).convert('RGB')):
        yield progress

if __name__ == "__main__":
    # Default input filename
    input_filename = "sample.jpg"

    # Check if a filename was passed as a command-line argument
    if len(sys.argv) > 1:
        input_filename = sys.argv[1]

    # scale the image to fit the screen (not 8k because of performance)
    rgb_values = scale_image(input_filename)
    
    compute_desaturation(rgb_values)