slice.py

from PIL import Image
import json
import numpy as np
import os
from time import sleep
import subprocess
import traceback


def read_time_values(filename):
    # # Order: R,G,B,UV
    # # each subarray represents the desaturation time needed for each pixel
    # time_array = [[],[],[],[]]

    # # parse the array 
    # # from [(R,G,B,UV)] to [[R], [G], [B], [UV]]
    # with open(filename, 'r') as file:
    #     # read out the R/G/B/UV values
    #     # result looks like [(10,20,20,30), (10,20,20,30) ... ]
    #     for line in file:
    #         # Split the line into parts, strip whitespace, and convert each to an integer
    #         parts = line.strip().split(',')
    #         assert(len(parts) == 4)
    #         for i in range(4):
    #             time_array[i].append(int(parts[i]))
    # return time_array
    
    with open(filename, 'r') as fin:
        return np.array(json.load(fin), dtype=np.int16)


# time_array: the time needed 
# interval: 
def create_bw_images(time_array_all, interval, output_folder, info_file_name="info.txt", order=(3,0,1,2)):
    # for frame counting purpose
    total_frame = 0
        
    d, height, width = time_array_all.shape
    
    assert(d == 4)
    info = {}
    info['interval'] = interval

    keys = ["R", "G", "B", "UV"]
    # iterate through R,G,B,UV
    # in the order of UV, R, G, B
    for color_index in order:
        time_array = time_array_all[color_index]
        max_frames = np.max(time_array) // interval + 1
        info[keys[color_index]] = int(max_frames)
        
        for frame in range(max_frames):
            total_frame += 1
            # Create an empty black image for the current frame
            img = Image.new('1', (width, height), "black")
            pixels = img.load()
            
            for i in range(height):
                for j in range(width):
                    if frame * interval < time_array[i][j]:
                        pixels[j, i] = 1
            
            # for i, duration in enumerate(time_array):
            #     # Determine if the current pixel should be white in this frame
            #     if frame * interval < duration:
            #         x, y = i % height, i // height
            #         if x < width and y < height:  # Check within image bounds
            #             pixels[x, y] = 1  # Set the pixel to white
            
            # Save the image
            img.transpose(Image.FLIP_TOP_BOTTOM)
            
            name = "0" * (5 - len(str(total_frame))) + str(total_frame)
            img.save(f"{output_folder}/{name}.png")
        
        print("%s: %s frames" % (keys[color_index], max_frames))
        
    with open("%s/%s" % (output_folder, info_file_name), 'w+') as fout:
        json.dump(info, fout)
        print("%s created in /%s" % (info_file_name, output_folder))

def run_slice(filename='out.tarr', interval=10, output_folder='sliced_pngs'):
    os.system("./clear_old_pngs.sh")
    time_values = read_time_values(filename)
    create_bw_images(time_values, interval, output_folder)
    print("Images have been generated based on the provided time values.")

def run_slice_tiled(canvasWidthCm, canvasHeightCm, screenWidthCm, screenHeightCm, screenWidthPx, screenHeightPx, tile_corners, tile_dims, method, frame_resolution, env_variables, tarr_file_name='out.tarr', info_file_name='info.json', output_folder='slice_out'):        
    os.system(f"./clear_old_pngs.sh {output_folder}")
    tarr = read_time_values(tarr_file_name)
    d, imgHeightPx, imgWidthPx = tarr.shape
    assert(d == 4)
                
    keys = ["R", "G", "B", "UV"]
    order = (3, 0, 1, 2)
    info = {}
    info['tiles'] = len(tile_corners)
    info['tile_corners'] = tile_corners
    info['tile_dims'] = tile_dims
    info['frames'] = []
    info['intervals'] = []
    
    imgAspectRatio = imgWidthPx / imgHeightPx
    canvasAspectRatio = canvasWidthCm / canvasHeightCm
    
    if (method == 0 and imgAspectRatio >= canvasAspectRatio) or (method == 1 and imgAspectRatio < canvasAspectRatio):
        imgHeightCm = canvasHeightCm
        imgWidthCm = canvasHeightCm * imgAspectRatio
    elif (method == 0 and imgAspectRatio < canvasAspectRatio) or (method == 1 and imgAspectRatio >= canvasAspectRatio):
        imgWidthCm = canvasWidthCm
        imgHeightCm = canvasWidthCm / imgAspectRatio
    elif (method == 2):
        imgWidthCm = canvasWidthCm
        imgHeightCm = canvasHeightCm
        
    imgStartHeightCm = (canvasHeightCm - imgHeightCm) / 2
    imgStartWidthCm = (canvasWidthCm - imgWidthCm) / 2
    
    def getImgPixel(widthCm, heightCm):
        return (imgWidthPx * ((widthCm - imgStartWidthCm) / imgWidthCm), 
                imgHeightPx * ((heightCm - imgStartHeightCm) / imgHeightCm))
    
    def getTarrVal(widthCm, heightCm, color_index, style='round'):
        if widthCm < 0 or widthCm >= canvasWidthCm or heightCm < 0 or heightCm >= canvasHeightCm:
            return 0
        i, j = getImgPixel(widthCm, heightCm)
        if style == 'round':
            i, j = round(i), round(j)
        if i < 0 or j < 0 or i >= imgWidthPx or j >= imgHeightPx:
            return 0
        elif color_index is None:
            return tarr[0:4, j, i]
        return tarr[color_index, j, i]
        
    def iter_tile(tileStartWidthCm, tileStartHeightCm, color_index=None):
        for i in range(screenHeightPx):
            for j in range(screenWidthPx):
                hCm = tileStartHeightCm + (i / screenHeightPx) * screenHeightCm
                wCm = tileStartWidthCm + (j / screenWidthPx) * screenWidthCm
                yield getTarrVal(wCm, hCm, color_index), j, i
        
    times = np.zeros(4)
    processes = []
    
    
    # START OS OPS
    
    try:
        
        for tile, (tileStartWidthCm, tileStartHeightCm) in enumerate(tile_corners):
            yield(0.0)
            
            max_frames_color = np.zeros(4, dtype=np.int16)
            for t, i, j in iter_tile(tileStartWidthCm, tileStartHeightCm):
                max_frames_color = np.maximum(max_frames_color, t)
                
            yield(0.05)
                
            intervals = np.ceil(max_frames_color / frame_resolution)
            for i in range(len(max_frames_color)):
                if intervals[i] == 0:
                    max_frames_color[i] = 0
                else: max_frames_color[i] = np.ceil(max_frames_color[i] / intervals[i])
            # max_frames_color = np.ceil(max_frames_color / intervals)
            inff = {}
            infi = {}
            for i, t in enumerate(max_frames_color):
                inff[keys[i]] = int(t)
                infi[keys[i]] = int(intervals[i])
            info['frames'].append(inff)
            info['intervals'].append(infi)
            sum_frames = np.sum(max_frames_color)
            total_frame = 0
            
            tile_name = "tile" + "0" * (3 - len(str(tile))) + str(tile)
            os.mkdir(f"{output_folder}/{tile_name}")
            
            for color_index in order:      
                os.mkdir(f"{output_folder}/{tile_name}/{keys[color_index]}")
                    
                for frame in range(max_frames_color[color_index]):
                    if (total_frame % max(1, sum_frames // 50) == 0):
                        yield(0.05 + (total_frame / sum_frames) * 0.95)
                        
                    img = Image.new('1', (screenWidthPx, screenHeightPx), "black")
                    pixels = img.load()
                    
                    for t, i, j in iter_tile(tileStartWidthCm, tileStartHeightCm, color_index):
                        if t > frame * intervals[color_index]:
                            pixels[i, screenHeightPx - 1 - j] = 1

                    # Save the image
                    name = "0" * (5 - len(str(frame))) + str(frame)
                    img.save(f"{output_folder}/{tile_name}/{keys[color_index]}/{name}.png")
                    
                    total_frame += 1
                    
                # if max_frames_color[color_index] > 0:
                #     folder = f"{output_folder}/{tile_name}/{keys[color_index]}"
                #     proc = subprocess.Popen([f'{env_variables["PATH_TO_SCRIPTS_SRC"]}/exec_scripts/build.sh', folder, str(intervals[color_index])], shell=False, env=env_variables)
                #     processes.append(proc)
                            
            times += max_frames_color * intervals
            yield tile, [int(times[i]) for i in order]
            
        yield(0.0)
        for i, p in enumerate(processes):
            p.wait()
            yield(((i + 1) / len(processes)) * 0.95)
            
        with open("%s/%s" % (output_folder, info_file_name), 'w+') as fout:
            json.dump(info, fout)
            print("%s created in /%s" % (info_file_name, output_folder))
            
        yield(1.0)
    except Exception as e:
        print(f"Failed: {e}")
        for p in processes:
            p.terminate()

def run_print(env_variables):
    processes = []
    keys = ['UV', 'R', 'G', 'B']
    
    try:
        yield 0.0
        proc = subprocess.Popen([f'{env_variables["PATH_TO_SCRIPTS_SRC"]}/exec_scripts/PI_move_files.sh'], shell=False, env=env_variables)
        processes.append(proc)
        proc.wait()
                
        with open(f"{env_variables['OUTPUT_FOLDER']}/info.json", 'r') as fin:
            info = json.load(fin)
            
        for i in range(info['tiles'])[3:]:
            width, height = info['tile_corners'][i]
            frames = info['frames'][i]
            intervals = info['intervals'][i]
            
            print(f"Moving robot to {width}, {height}")
            proc = subprocess.Popen([f'{env_variables["PATH_TO_SCRIPTS_SRC"]}/exec_scripts/PI_run_util.sh', str(width), str(height)], shell=False, env=env_variables)
            processes.append(proc)
            proc.wait()
            
            input(f"Press ENTER when screen is positioned for frame {i+1}: ")
            
            for k in keys[3:]:
                
                yield 0.2 + 0.8 * (i * len(keys) + keys.index(k)) / (info['tiles'] * len(keys))
                
                if frames[k] != 0:
                    
                    led_time = frames[k] * intervals[k]
                    total_time = led_time + frames[k] * int(env_variables['SWITCH_DELAY'])
                    
                    print(f"Printing channel {k} on frame {i+1}")
                    
                    proc = subprocess.Popen([f'{env_variables["PATH_TO_SCRIPTS_SRC"]}/exec_scripts/PI_start_print.sh', str(i), str(k)], shell=False, env=env_variables)
                    processes.append(proc)
                    proc.wait()
                    
                    print(f"Shining channel {k} for {total_time}s")
                    
                    proc = subprocess.Popen([f'{env_variables["PATH_TO_SCRIPTS_SRC"]}/exec_scripts/PI_set_backlight.sh', str(k), str(total_time)], shell=False, env=env_variables)
                    processes.append(proc)
                    proc.wait()
                    
                    #sleep(int(env_variables['START_DELAY']))
                    
                    #sleep(total_time + int(env_variables['START_DELAY']))
                    
                    input("Press ENTER when ready for next channel: ")
                    
                    proc = subprocess.Popen([f'{env_variables["PATH_TO_SCRIPTS_SRC"]}/exec_scripts/PI_stop_print.sh'], shell=False, env=env_variables)
                    processes.append(proc)
                    proc.wait()
                    
                    proc = subprocess.Popen([f'{env_variables["PATH_TO_SCRIPTS_SRC"]}/exec_scripts/PI_set_backlight.sh', 'O', '0'], shell=False, env=env_variables)
                    processes.append(proc)
                    proc.wait()
                    
                    sleep(3)
                    
                    print(f"Channel {k} on frame {i+1} finished")
                            
        yield 1.0
        
    except Exception as e:
        print(f"Running print failed with {e}")
        print(traceback.format_exc())
        for proc in processes:
            proc.terminate()
    
if __name__ == "__main__":
    # Parameters
    filename = 'out.tarr'  # Name of the file with time values
    interval = 10  # Interval for slicing time values
    # width, height = 2048, 1080  # Dimensions of the generated images
    output_folder = 'sliced_pngs'  # Folder to save the generated images

    # Process
    time_values = read_time_values(filename)
    create_bw_images(time_values, interval, output_folder)
    print("Images have been generated based on the provided time values.")