Commit 547da3ec authored by Vít Starý Novotný's avatar Vít Starý Novotný
Browse files

Update preprocess.apply_super_resolution()

parent 79cb84a5

ahisto_ocr/default.ini

+1 −0
Original line number Diff line number Diff line 
[preprocessing]

nvidia_visible_devices = all

model = anime_style_art_rgb

noise_level = 3

ahisto_ocr/preprocessing.py

+30 −7
Original line number Diff line number Diff line
@@ -16,6 +16,7 @@ CONFIG = _CONFIG['preprocessing'] 


def apply_super_resolution(client, input_volume, postprocessing_volume) -> None:

    LOGGER.info('Pre-processing the input images using super-resolution')



    volumes = {

        input_volume.name: {

            'bind': '/input',
@@ -26,17 +27,39 @@ def apply_super_resolution(client, input_volume, postprocessing_volume) -> None: 
            'mode': 'rw',

        },

    }



    command = [

        'parallel', '--halt', 'now,fail=1', '--jobs', '1', '--',

        'cp {} /output/{/}', '::::', '/input/list.txt',

    ]

    run_docker_container(client, 'ahisto/ocr-eval', command=command, volumes=volumes)



    with create_temporary_docker_container(client, 'ahisto/empty', command='cmd',

                                           volumes=volumes) as container:

        filenames = [

            '{}_noise_scale.png'.format(Path(filename).with_suffix(''))

            for filename

            in extract_text_file_from_container(container, '/input/list.txt').split('\n')

        ]

        add_text_file_to_container(container, '/output/list.txt', '\n'.join(filenames))

        text = extract_text_file_from_container(container, '/input/list.txt')

        add_text_file_to_container(container, '/output/list.txt', text)



    volumes = {

        postprocessing_volume.name: {

            'bind': '/input',

            'mode': 'rw',

        },

    }



    command = [

        'th', '/root/waifu2x/waifu2x.lua', '-force_cudnn', '1',

        '-model_dir', f'/root/waifu2x/models/{CONFIG["model"]}', '-m', 'noise_scale',

        '-scale', '2', '-noise_level', str(CONFIG.getint('noise_level')), '-l', '/input/list.txt',

    ]

    run_docker_container(client, 'ahisto/waifu2x', gpu=True, command=command, volumes=volumes, working_dir='/output')
 
    run_docker_container(client, 'ahisto/waifu2x', runtime='nvidia',

                         environment={'NVIDIA_VISIBLE_DEVICES': CONFIG['nvidia_visible_devices']},

                         command=command, volumes=volumes)



    with create_temporary_docker_container(client, 'ahisto/empty', command='cmd',

                                           volumes=volumes) as container:

        filenames = [

            '{}_noise_scale.png'.format(Path(filename).with_suffix(''))

            for filename

            in extract_text_file_from_container(container, '/input/list.txt').split('\n')

        ]

        add_text_file_to_container(container, '/input/list.txt', '\n'.join(filenames))

ahisto_ocr/util.py

+2 −6
Original line number Diff line number Diff line
@@ -30,9 +30,7 @@ def create_temporary_docker_volume(client, *args, **kwargs) -> Iterator: 




@contextmanager

def create_temporary_docker_container(client, *args, gpu: bool = False, **kwargs) -> Iterator:

    if gpu:

        kwargs = {**kwargs, **{'device_requests': [DeviceRequest(count=-1, capabilities=[['gpu']])]}}

def create_temporary_docker_container(client, *args, **kwargs) -> Iterator:

    container = client.containers.create(*args, **kwargs)

    LOGGER.debug(f'Created temporary Docker container {container.short_id}')

    try:
@@ -42,10 +40,8 @@ def create_temporary_docker_container(client, *args, gpu: bool = False, **kwargs 
        LOGGER.debug(f'Deleted temporary Docker container {container.short_id}')





def run_docker_container(client, *args, gpu: bool = False, quiet: bool = True, **kwargs) -> None:

def run_docker_container(client, *args, quiet: bool = True, **kwargs) -> None:

    kwargs = {**{'remove': True, 'stdout': True, 'stderr': True}, **kwargs}

    if gpu:

        kwargs = {**kwargs, **{'device_requests': [DeviceRequest(count=-1, capabilities=[['gpu']])]}}

    try:

        output = client.containers.run(*args, **kwargs)

        if not quiet: