Deoldify, как использовать локальную модель?

Рейтинг: 0Ответов: 0Опубликовано: 04.07.2023

Вдохновившись этим вопросом, поставил себе локально torch и deoldify.
Пришлось повозиться с установкой, т.к. привык к коллабам, а не локальному запуску, но работает отлично.

По библиотеке есть их сайт ссылающийся на коллаб, а также github, описывающий локальный запуск и ссылающийся на .ipynb, которых нет в репозитории, но есть другие. Видимо документация немного устарела.
Не важно, по тем что есть в репозитории код немного отличается, в основном предлагая использовать фотографии на локальном диске, а не url.

В коде есть следующая строчка:

colorizer = get_image_colorizer()

В collab она отрабатывает как задумано. Если нет скачанной модели, то ругается. В jupyter-notebook всё немного иначе. Даже если по пути ./models есть модель, функция всё равно скачивает другую в кэш:

Downloading: "https://download.pytorch.org/models/resnet101-63fe2227.pth" to C:\Users\User/.cache\torch\hub\checkpoints\resnet101-63fe2227.pth
 33%|█████████████████████████▌                                                    | 56.0M/171M [00:09<00:18, 6.34MB/s]

При этом, стоит отметить, что при ключе artistic = False скачивается модель размером 171M, хотя по ссылкам на GH ColorizeStable_gen.pth весит 853M, т.е. что он скачивает не совсем понятно. Я скачал все модели и поместил в соответствующую папку в root проекта.

Конечно же, я посмотрел что из себя представляет функция:

<function deoldify.visualize.get_image_colorizer(root_folder: pathlib.Path = WindowsPath('.'), render_factor: int = 35, artistic: bool = True) -> deoldify.visualize.ModelImageVisualizer>

И попробовал использовать ключ root_folder c from pathlib import Path используя путь к root проекта, но это не помогает. Каждый раз функция скачивает модель в .cache, не создаёт папку models в User root и игнорирует содержимое папки проекта…

Вопрос

Как использовать локально загруженные модели.
Это тем более критично, потому что скачиваемые автоматически явно меньше чем загруженные с GH.


Поискал по самому репозиторию упоминание get_image_colorizer, естественно нашлись .ipynb, а также по пути "C:\Users\User\Documents\GitHub\DeOldify\deoldify\visualize.py":

def get_image_colorizer(
    root_folder: Path = Path('./'), render_factor: int = 35, artistic: bool = True
) -> ModelImageVisualizer:
    if artistic:
        return get_artistic_image_colorizer(root_folder=root_folder, render_factor=render_factor)
    else:
        return get_stable_image_colorizer(root_folder=root_folder, render_factor=render_factor)

def get_stable_image_colorizer(
    root_folder: Path = Path('./'),
    weights_name: str = 'ColorizeStable_gen',
    results_dir='result_images',
    render_factor: int = 35
) -> ModelImageVisualizer:
    learn = gen_inference_wide(root_folder=root_folder, weights_name=weights_name)
    filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
    vis = ModelImageVisualizer(filtr, results_dir=results_dir)
    return vis

Что именно начинает загрузку, и как использовать локальную модель — мне пока не понятно. В ModelImageVisualizer нет никаких признаков скачивания модели.
Да, я пока не пробовал запускать дебаг, возможно ответ на поверхности. Если нет, отвечу сам себе.

Ответы

Ответов пока нет.