stopi2/frame_opencv.py

import cv2
import gettext
from itertools import count
import os
# Needed for utf-8 text
from PIL import ImageFont, ImageDraw, Image
from send2trash import send2trash
import signal
import sys
import subprocess
import time
from timeit import default_timer as timer
import tomllib
import numpy as np
import serialutils

running_from_folder = os.path.realpath(__file__)
alphabet = ['A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z']


# l10n
LOCALE = os.getenv('LANG', 'en_EN')
_ = gettext.translation('template', localedir='locales', languages=[LOCALE]).gettext

# Config 
# defaults
project_settings_defaults = {
    'cam_is_showmewebcam': False,
    'use_date_for_folder': False,
    'file_extension':'png',
    'projects_folder': '',
    'onion_skin_onstartup' : False,
    'onionskin_alpha_default' : 0.4,
    'fullscreen_bool' : True,
    'screen_w' : 1920,
    'screen_h' : 1080,
    'framerate' : 16,
    'ffmpeg_path' : None,
    'export_options' : 'scale=1920:-1,crop=1920:1080:0:102',
}

camera_current_settings_defaults = {
    'cam_w' : 800,
    'cam_h' : 600,
    'vflip' : 0,
    'hflip' : 0,
    'auto_exposure' : 0,
    'white_balance_auto_preset' : 0,
}

# Load from file
config_locations = ["./", "~/.", "~/.config/"]
config_found_msg = _("No configuration file found, using defaults.")
project_settings = project_settings_defaults
camera_current_settings = camera_current_settings_defaults
for location in config_locations:
    # Optional config files, ~ is expanded to $HOME on *nix, %USERPROFILE% on windows
    if os.path.exists( os.path.expanduser(os.path.join(location, 'config.toml'))):
        with open(os.path.expanduser(location + 'config.toml'), 'rb') as config_file:
            project_settings = tomllib.load(config_file)
            if 'CAMERA' in project_settings:
                camera_settings = project_settings['CAMERA']
            if 'DEFAULT' in project_settings:
                project_settings = project_settings['DEFAULT']
            config_found_msg = _("Found configuration file in {}").format(os.path.expanduser(location))
print(config_found_msg)

camera_current_settings = {
        'auto_exposure': dict(min=0, max=1, default=camera_settings['auto_exposure'], value=camera_settings['auto_exposure']),
        'white_balance_auto_preset': dict(min=0, max=9, default=camera_settings['white_balance_auto_preset'], value=camera_settings['white_balance_auto_preset']),
        'horizontal_flip': dict(min=0, max=1, default=camera_settings['hflip'], value=camera_settings['hflip']),
        'vertical_flip': dict(min=0, max=1, default=camera_settings['vflip'], value=camera_settings['vflip']),
        'video_bitrate': dict(min=25000000, max=25000000, default=camera_settings['video_bitrate'], value=camera_settings['video_bitrate']),
    }


def apply_cam_setting(cam_settings:dict, to_set:list=None):
    cmd = 'v4l2-ctl -d /dev/video0'
    args = []
    for setting in cam_settings:
        if to_set is None:
            # Apply defaults
            cam_settings[setting]['value'] = cam_settings[setting]['default']
            cmd += ' -c {}={}'
            args.append(setting)
            args.append(cam_settings[setting]['value'])
        else:
            # Increment settings in to_set
            for setting in to_set:
                if setting in cam_settings:
                    if cam_settings[setting]['value']+1 in range(cam_settings[setting]['min'],cam_settings[setting]['max']+1):
                        cam_settings[setting]['value'] += 1
                    else:
                        cam_settings[setting]['value'] = cam_settings[setting]['min']
                    cmd += ' -c {}={}'
                    args.append(setting)
                    args.append(cam_settings[setting]['value'])
                else:
                    print(_("Unknown setting!"))
            break
    if project_settings['cam_is_showmewebcam']:
        serialutils.send_serial_cmd(serialutils.find_cam_port(), cmd.format(*args))
    else:
        # TODO: v4l2 support
        print(_("Camera function not supported."))
    return cam_settings


def generate_text_image(text:str, screen_w, screen_h, bullets=False):
    text_image = Image.new('RGB',
                            (screen_w, screen_h),
                            (0,0,0)
                          )
    text_image_draw = ImageDraw.Draw(text_image)
    if text is not None:
        font = ImageFont.truetype("Tuffy_Bold.ttf", (screen_w/32))
        lines = text.split('\n')
        longest_line = lines[0]
        for line in lines:
            if len(line) > len(longest_line):
                longest_line = line
        font_len = font.getlength(lines[lines.index(longest_line)])
        text_image_draw.multiline_text((screen_w/2 - font_len/2, screen_h/3 ),
                                        text,
                                        fill=(255, 255, 255),
                                        font=font,
                                        align='center',
                                        spacing=20
                                        )
    if bullets:
        dot_radius = screen_w/24
        x_unit = (screen_w/32)
        y_unit = (screen_h/32)
        green_dot = (x_unit*14, y_unit*24)
        red_dot   = (green_dot[0]+x_unit*4, green_dot[1])
        # Green dot
        text_image_draw.circle(green_dot,
                         dot_radius,
                         fill=(0,255,0),
                         outline=None,
                         width=1
                         )
        text_image_draw.text(green_dot,
                            _("Yes"),
                            fill=(0, 0, 0),
                            font=font,
                            anchor="mm",
                            spacing=20
                            )
        # Red dot
        text_image_draw.circle(red_dot,
                         dot_radius,
                         fill=(0,0,255),
                         outline=None,
                         width=1
                         )
        text_image_draw.text(red_dot,
                            _("No"),
                            fill=(0, 0, 0),
                            font=font,
                            anchor="mm",
                            spacing=20
                            )
    text_image = np.array(text_image)
    return text_image


def askyesno(text:str):
    blank = generate_text_image(text, project_settings['screen_w'], project_settings['screen_h'], bullets=True)
    cv2.imshow("StopiCV", blank)
    # Wait for input to continue
    answer = cv2.waitKey(0)
    # Space pressed == yes
    if answer%256 == 32 or answer%256 == 48 :
        return True
    # Any other key == no
    else:
        return False


def find_letter_after(letter:str, date=False):
    if letter in alphabet and alphabet.index(letter) < len(alphabet) - 1 and not date:
        letter = alphabet[alphabet.index(letter) + 1]
    else:
        # Create folder with date
        year,mon,day,hour,minute,sec,wd,yd,dst = time.localtime()
        letter = '{}-{}-{}_{}-{}-{}'.format(year,mon,day,hour,minute,sec)
    return letter


def get_projects_folder():
    if len(projects_folder):
        project_folder = projects_folder
    else:
        # Get user folder
        project_folder = os.path.expanduser('~')
        # If a project folder is defined in settings, use it
        if project_settings['projects_folder'] != '':
            subfolder = project_settings['projects_folder']
        else:
            # If it doesn't exist, use a default name
            subfolder = 'Stopmotion Projects'
        project_folder = os.path.join(project_folder, subfolder)
    # Create folder if it doesn't exist
    if os.path.exists(project_folder) == False:
        os.mkdir(project_folder)
    else:
        if not os.path.isdir(project_folder):
            # If file exists but is not a folder, can't create it, abort
            return False
    return project_folder


def get_session_folder():
    global next_letter
    project_folder = get_projects_folder()
    if project_folder:
        sessions_list = []
        dir_list = os.listdir(project_folder)
        # Filter folders with name only one char long
        for folder in dir_list:
            if len(folder) == 1 and folder in alphabet:
                sessions_list.append(folder)
        # If folders exist, find last folder in alphabetical order
        if len(sessions_list):
            sessions_list.sort()
            last_letter = sessions_list[-1]
            # By default, find next letter for a new session
            next_letter = find_letter_after(last_letter, project_settings['use_date_for_folder'])
            if next_letter is False:
                return False
            # A previous session folder was found; ask the user if they wish to resume session
            if not project_settings['use_date_for_folder']:
                resume_session = askyesno(_("A previous session was found in\n {},\n resume shooting ?").format(os.path.join(project_folder, last_letter)))
                # ~ resume_session = tk.messagebox.askyesno(_("Resume session?"), _("A previous session was found in {}, resume shooting ?").format(os.path.join(project_folder, last_letter)))
                if resume_session:
                    next_letter = last_letter
        else:
            if not project_settings['use_date_for_folder']:
                next_letter = 'A'
            else:
                next_letter = find_letter_after('A', project_settings['use_date_for_folder'])
        if os.path.exists(os.path.join(project_folder, next_letter)) is False:
            os.mkdir(os.path.join(project_folder, next_letter))
        print(_("Using {} as session folder.").format(os.path.join(project_folder, next_letter)))
        return os.path.join(project_folder, next_letter)
    return False


def get_frames_list(folder:str):
    # Get JPG files list in current directory
    existing_animation_files = img_list
    file_list = os.listdir(folder)
    for file in file_list:
        if (file.startswith(project_letter) and file.endswith(project_settings['file_extension'])):
            if file not in existing_animation_files:
                existing_animation_files.append(file)
    if len(existing_animation_files) == 0:
        # If no images were found, return fake name set to -001 to init file count to 000
        return ["{}.{:04d}.{}".format(next_letter, -1, project_settings['file_extension'])]
    existing_animation_files.sort()
    return existing_animation_files


def get_last_frame(folder:str):
    # Refresh file list
    existing_animation_files = get_frames_list(folder)
    # Get last file
    # Filename pattern is A.0001.JPG
    return existing_animation_files[-1].split('.')


def get_onionskin_frame(folder:str, index=None):
    prev_image = get_last_frame(folder)
    prev_image = '.'.join(prev_image)
    if os.path.exists( os.path.expanduser(os.path.join(savepath, prev_image))):
        frm = cv2.imread(os.path.join(savepath, prev_image))
        frm = cv2.resize(frm, (project_settings['screen_w'], project_settings['screen_h']))
    # Imge does not exist, load blank image
    else:
        frm = blank_image
    return frm


def return_next_frame_number(last_frame_name):
    prefix, filecount, ext = last_frame_name
    filename = '.{:04d}.'.format(int(filecount)+1)
    return prefix + filename + ext


def update_image(img_list, img_index):
    if len(img_list) == 0:
        return 0
    img_filename = img_list[img_index]
    if os.path.exists( os.path.expanduser(os.path.join(savepath, img_filename))):
        img = cv2.imread(os.path.join(savepath, img_filename))
        img = cv2.resize(img, (project_settings['screen_w'], project_settings['screen_h']))
    else:
        img = blank_image
    return img

def next_frame(img_index, loop=True):
    img_index = check_range(img_index+1, loop)
    return img_index, update_image(img_list, img_index)
    

def previous_frame(img_index):
    img_index = check_range(img_index-1)
    return img_index, update_image(img_list, img_index)


def last_frame(img_index):
    img_index = len(img_list)-1
    return img_index, update_image(img_list, img_index)


def first_frame(img_index):
    img_index = 0
    return img_index, update_image(img_list, img_index)


def clean_img_list(folder_path):
    # Check file in dict exists, else remove it
    file_list = os.listdir(folder_path)
    # Iterate over copy of dict to avoid OOR error
    img_list_copy = img_list
    for file in img_list_copy:
        if file not in file_list:
            img_list.remove(file)


def check_range(x, loop=True):
    if x < 0:
        if loop:
            return len(img_list)-1
        else:
            return 0
    elif x >= len(img_list)-1:
        if loop:
            return 0
        else:
            return len(img_list)-1
    else:
        return x


def batch_rename(folder:str):        
    # initialize counter to 0
    frame_list = get_frames_list(folder)
    counter = (".%04i." % x for x in count(0))
    for i in frame_list:
        if os.path.exists(os.path.join(folder, i)):
            os.rename(os.path.join(folder, i), os.path.join(folder, "{}{}{}".format(project_letter, next(counter), project_settings['file_extension'])))
        else:
            print(_("{} does not exist").format(str(i)))
    return get_frames_list(folder)


def offset_dictvalues(from_index=0):
    dict_copy = dict(img_list)
    for i in range(from_index, len(dict_copy)):
        if i < len(img_list)-1:
            img_list[list(img_list.keys())[i]] = list(img_list.values())[i+1]
        else:
            img_list[list(img_list.keys())[i]] = None


def remove_frame(img_list, img_index):
    if len(img_list):
        folder_path = os.path.realpath(savepath)
        frame_name = img_list[img_index]
        # ~ frame_path = os.path.realpath(frame_name)
        frame_path = os.path.join(folder_path, frame_name)
        if not os.path.exists(frame_path):
            return img_list, img_index, blank_image
        print(_("Removing {}").format(frame_path))
        # trash file
        send2trash(frame_path)
        # remove entry from dict
        img_list.remove(frame_name)
        # rename files and get new list
        img_list = batch_rename(folder_path)
        clean_img_list(folder_path)
        # update index if possible
        img_index = check_range(img_index, False)
        # update display
        return img_list, img_index, update_image(img_list, img_index)
    else:
        return img_list, 0, blank_image


def testDevice(source):
    cap = cv2.VideoCapture(source) 
    if cap is None or not cap.isOpened():
        print(_("Warning: unable to open video source: {}").format(source))
        return False
    cap.release()
    return True


def signal_handler(sig, frame):
    global ctrlc_pressed
    ctrlc_pressed = True


def parse_export_options(options:str, vflip:int=0, hflip:int=0):
    if vflip:
        options += ',vflip'
    if hflip:
        options += ',hflip'
    return options


def export_animation(input_filename, export_filename):
    input_format, framerate = input_options
    if project_settings['ffmpeg_path'] is None:
        return False
    ffmpeg_process = subprocess.Popen([
        project_settings['ffmpeg_path'],
        '-v','quiet',
        '-y',
        '-f', input_format, 
        '-r', framerate,
        '-i', input_filename,
        '-vf', output_options,
        # Fix for Quicktime
        '-pix_fmt', 'yuv420p',
        '-vcodec', 'libx264',
        export_filename,
    ])
    return ffmpeg_process


def main(args):

    global onionskin, liveview_only, playback, loop_playback, playhead, index, img_list, first_playback, camera_current_settings
    
    if not testDevice(0):
        print(_("No camera device found. Exiting..."))
        return 1
    cam = cv2.VideoCapture(0)
    cam.set(cv2.CAP_PROP_FRAME_WIDTH, camera_settings['cam_w'])
    cam.set(cv2.CAP_PROP_FRAME_HEIGHT, camera_settings['cam_h'])

    frame = get_onionskin_frame(savepath, index)
   
    if project_settings['cam_is_showmewebcam']:
        # Apply config camera settings 
        camera_current_settings = apply_cam_setting(camera_current_settings)
        time.sleep(.5)
    
    loop_delta = 0
    while True:
        start = timer()
        if playback:
            if onionskin:
                onionskin = False
                onionskin_was_on = True
            if first_playback:
                playhead = 0
                first_playback = False
            # Play all frames
            if playhead < len(img_list)-1:
                playhead, img = next_frame(playhead, loop_playback)
                cv2.imshow("StopiCV", img)
                # Calculate framerate according to loop execution time
                frame_interval = 1.0/project_settings['framerate']-loop_delta
                if frame_interval < 0:
                    frame_interval = 0
                time.sleep(frame_interval)
            else:
                playhead = index
                img = update_image(img_list, index)
                first_playback = True
                playback = False
                # Restore onionskin
                if 'onionskin_was_on' in locals():
                    onionskin = True
            loop_playback = False
        
        if liveview_only:
            # ~ onionskin = False
            ret, overlay = cam.read()
            if not ret:
                print(_("Failed to grab frame."))
                break
            cv2.imshow("StopiCV", overlay)
        # ~ else:
             # ~ onionskin = True
                
        if onionskin:
            ret, overlay = cam.read()
            og_frame = overlay.copy()
            # Resize preview
            overlay = cv2.resize(overlay, (project_settings['screen_w'], project_settings['screen_h']))
            # Apply onionskin
            alpha = project_settings['onionskin_alpha_default']
            beta = (1.0 - alpha)
            overlay = cv2.addWeighted(frame, alpha, overlay, beta, 0)

            if not ret:
                print(_("Failed to grab frame."))
                break
            cv2.imshow("StopiCV", overlay)
        
        if not playback and not onionskin and not liveview_only:
            cv2.imshow("StopiCV", frame)

        k = cv2.waitKey(1)
        # TODO : Show liveview only
        # Key l / 5
        if (k%256 == 108) or (k%256 == 53):
            # Toggle liveview
            liveview_only = not liveview_only
            onionskin = not onionskin
        # Key o / 9
        if (k%256 == 111) or (k%256 == 47):
            # Toggle onionskin
            onionskin = not onionskin
            liveview_only = False
        # Key w / 7 - cycle wb
        if (k%256 == 119) or (k%256 == 55):
            camera_current_settings = apply_cam_setting(camera_current_settings, ['white_balance_auto_preset'])
        # Key x / 1 - cycle exposure
        if (k%256 == 120) or (k%256 == 49):
            camera_current_settings = apply_cam_setting(camera_current_settings, ['auto_exposure'])
        # Key f / 3 - flip image
        if (k%256 == 102) or (k%256 == 51):
            camera_current_settings = apply_cam_setting(camera_current_settings, ['vertical_flip','horizontal_flip'])
        # Key up
        elif (k%256 == 82) or (k%256 == 56):
            if len(img_list):   
                if playback:
                    playback = False
            index, frame = last_frame(index)
        # Key down
        elif (k%256 == 84) or (k%256 == 50):
            if len(img_list):   
                if playback:
                    playback = False
            index, frame = first_frame(index)
        # Key left
        elif (k%256 == 81) or (k%256 == 52):
            # Displau previous frame
            if len(img_list):
                if playback:
                    playback = False
                index, frame = previous_frame(index)
        # Key right
        elif (k%256 == 83) or (k%256 == 54):
            # Displau next frame
            if len(img_list):
                if playback:
                    playback = False
                index, frame = next_frame(index)
        # Key r / keypad 9 - reset wb,exp
        elif (k%256 == 114) or (k%256 == 57):
            print(_("Reset camera settings"))
            camera_current_settings = apply_cam_setting(camera_current_settings)
        # Key e / keypad *
        elif (k%256 == 101) or (k%256 == 42):
            print(_("Export"))
            ffmpeg_process = export_animation(input_filename, export_filename)
        # Key Return
        elif (k%256 == 13):
            print(_("Playback"))
            playhead = index
            loop_playback = True
            playback = not playback
        # Key remove frame - backspace, del, numpad_minus
        elif (k%256 == 8) or (k%256 == 45) or (k == 255):
            # Remove frame
            print(_("Remove frame"))
            img_list, index, frame = remove_frame(img_list, index)
        # Quit app
        elif k%256 == 27:
            # ESC pressed
            print(_("Escape hit, exiting..."))
            break
        elif ctrlc_pressed:
            print(_("Ctrl-C hit, exiting..."))
            break
        elif cv2.getWindowProperty("StopiCV", cv2.WND_PROP_VISIBLE) < 1:
            print(_("Window was closed, exiting..."))
            break
        # Take pic
        # SPACE or numpad 0 pressed
        elif (k%256 == 32) or (k%256 == 48):
            img_name = return_next_frame_number(get_last_frame(savepath))
            img_path = os.path.join(savepath, img_name)
            cv2.imwrite(img_path, og_frame)
            print(_("File {} written.").format(img_path))
            if len(img_list) and (img_list[index] == '{letter}.-001.{ext}'.format(letter=project_letter, ext=project_settings['file_extension'])):
                img_list[index] = img_name
            else:
                index += 1
            frame = get_onionskin_frame(savepath, index)
        # REMOVE : Debug print keycode
        elif k==-1:  # normally -1 returned,so don't print it
            continue
        else:
            print(k) # else print its value
        end = timer()
        loop_delta = end - start

    if 'ffmpeg_process' in locals():
         if ffmpeg_process.poll() is None:
            print(_("Ffmpeg is still running.\n Waiting for task to complete."))
            msg = generate_text_image(_("Ffmpeg is still running.\n Waiting for task to complete."),
                                project_settings['screen_w'], project_settings['screen_h']
                                )
            cv2.imshow("StopiCV", msg)
            # Force window refresh
            cv2.pollKey()
            try:
                ffmpeg_process.wait(timeout=20)
            except:
                print(_("Terminating running process..."))
                ffmpeg_process.terminate()
    cam.release()
    cv2.destroyAllWindows()

cv2.namedWindow("StopiCV", cv2.WINDOW_GUI_NORMAL)
cv2.setWindowProperty("StopiCV", cv2.WND_PROP_OPENGL, cv2.WINDOW_OPENGL)
cv2.setWindowProperty("StopiCV", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
cv2.setWindowProperty("StopiCV", cv2.WND_PROP_ASPECT_RATIO, cv2.WINDOW_KEEPRATIO)

ctrlc_pressed = False
projects_folder = project_settings['projects_folder']
next_letter = 'Z'
img_list = []
savepath = get_session_folder()
onionskin = project_settings['onion_skin_onstartup']
liveview_only = False
playback = False
first_playback = True
playhead = 0
loop_playback = True
blank_image = generate_text_image(_("No images yet! Start shooting..."), project_settings['screen_w'], project_settings['screen_h'])

if len(savepath):
    project_letter = savepath.split(os.sep)[-1]
else:
    project_letter = 'A'
img_list = get_frames_list(savepath)
index = len(img_list)-1

# Export settings
input_filename = "{folder}{sep}{letter}.%04d.{ext}".format(folder=savepath, sep=os.sep, letter=project_letter, ext=project_settings['file_extension'])
input_options = ["image2", str(project_settings['framerate'])]
# ~ output_filename = "{folder}{sep}{filename}.mp4".format(folder=projects_folder, sep=os.sep, filename=savepath.split(os.sep)[-1])
output_filename = "{filename}.mp4".format(filename=project_letter)
output_options = project_settings['export_options']

export_filename = os.path.join(savepath, output_filename)

if __name__ == '__main__':
    signal.signal(signal.SIGINT, signal_handler)
    sys.exit(main(sys.argv[1:]))