Справка по API: ML Компонент

Этот раздел описывает внутренние модули, отвечающие за алгоритмы обработки изображений.

Основной пайплайн

Этот модуль является точкой входа для запуска процесса восстановления изображения.

denoise_bot.ML_component.main_model.create_mask_from_damaged(damaged_image: ndarray, threshold: int = 10) ndarray[исходный код]

Отвечает за создание маски для матрицы.

Создает маску известных пикселей на основе поврежденного изображения (работает с NumPy массивами). Пиксель считается известным, если он не является почти черным.

Параметры:
  • damaged_image (np.ndarray) – Входное изображение в формате NumPy (H, W, C).

  • threshold (int) – Порог яркости для определения черного пикселя (в диапазоне 0-255).

Результат:

Двумерная булева маска, где True означает известный (неповрежденный) пиксель.

Тип результата:

np.ndarray

denoise_bot.ML_component.main_model.run_inpainting_pipeline(damaged_image_source: str | IO[bytes], output_dir: str, max_iters: int = 250, original_image_source: str | IO[bytes] | None = None, use_gpu: bool = True) ndarray | None[исходный код]

Главная функция для запуска всего процесса восстановления изображения (inpainting).

Оркестрирует загрузку данных, выбор вычислительного бэкенда (CPU/GPU), запуск алгоритма ADMM для матричного дозаполнения и сохранение результатов. Может принимать на вход как пути к файлам, так и байтовые потоки изображений.

Параметры:
  • damaged_image_source (Union[str, IO[bytes]]) – Источник поврежденного изображения. Может быть строкой с путем к файлу или файлоподобным объектом с байтами (например, io.BytesIO из Telegram).

  • output_dir (str) – Путь к папке, в которую будут сохранены все результаты (восстановленное изображение, графики). Папка будет создана, если не существует.

  • max_iters (int, optional) – Максимальное количество итераций для алгоритма ADMM. По умолчанию 250.

  • original_image_source (Optional[Union[str, IO[bytes]]], optional) – Источник оригинального (неповрежденного) изображения для создания сравнительной визуализации. Если не указан (None), сравнительный график создаваться не будет. По умолчанию None.

  • use_gpu (bool, optional) – Флаг для использования GPU. Если True, будет предпринята попытка использовать CuPy для вычислений. При его отсутствии или ошибке произойдет автоматический откат к NumPy (CPU). По умолчанию True.

Результат:

Восстановленное изображение в виде NumPy массива (на CPU) в случае успешного выполнения. Возвращает None, если в процессе произошла ошибка.

Тип результата:

Optional[np.ndarray]

Ядро алгоритма ADMM

Здесь реализованы математические шаги алгоритма матричного дозаполнения.

denoise_bot.ML_component.admm_core.MC_ADMM(Y: ArrayLike, mask: ArrayLike, tol: float, max_iters: int, r: float, backend: BackendModule) Tuple[ArrayLike, List[float]][исходный код]

Дозаполняет матрицу с помощью ADMM.

Решает задачу матричного дозаполнения с помощью ADMM, управляя итерационным процессом.

Параметры:
  • Y (ArrayLike) – Исходная матрица (канал изображения) с пропусками.

  • mask (ArrayLike) – Булева маска, где True - известные элементы.

  • tol (float) – Порог для критерия остановки (изменение ядерной нормы).

  • max_iters (int) – Максимальное количество итераций.

  • r (float) – Параметр rho для ADMM.

  • backend (BackendModule) – Вычислительный бэкенд для всех операций (модуль numpy или cupy).

Результат:

Кортеж, содержащий:
  • Восстановленную матрицу X (на том же бэкенде, что и входные данные).

  • Историю значений ядерной нормы (список чисел float).

Тип результата:

Tuple[ArrayLike, List[float]]

denoise_bot.ML_component.admm_core.MC_update_X(Z: ArrayLike, lambda_: ArrayLike, r: float, np: BackendModule) ArrayLike[исходный код]

Обновление матрицы X.

Шаг обновления низкоранговой матрицы X через пороговое сжатие сингулярных чисел (SVT). Это соответствует решению задачи: argmin_X ||X||_* + (r/2)||X - Z + U||^2_F

Параметры:
  • Z (ArrayLike) – Текущая оценка матрицы, удовлетворяющей ограничениям.

  • lambda (ArrayLike) – Двойственная переменная (масштабированная).

  • r (float) – Параметр rho для ADMM, контролирующий штраф.

  • np (BackendModule) – Вычислительный бэкенд (модуль numpy или cupy).

Результат:

Обновленная низкоранговая матрица X

того же типа, что и входные.

Тип результата:

ArrayLike

denoise_bot.ML_component.admm_core.MC_update_Z(X: ArrayLike, lambda_: ArrayLike, r: float, Y: ArrayLike, mask: ArrayLike) ArrayLike[исходный код]

Обновление матрицы Z.

Шаг обновления матрицы Z, удовлетворяющей ограничениям на известные пиксели. Это соответствует проекции на множество матриц, совпадающих с Y на маске E.

Параметры:
  • X (ArrayLike) – Текущая оценка низкоранговой матрицы.

  • lambda (ArrayLike) – Двойственная переменная.

  • r (float) – Параметр rho для ADMM.

  • Y (ArrayLike) – Исходная матрица с пропусками (используется для известных значений).

  • mask (ArrayLike) – Булева маска, где True - известные элементы.

Результат:

Обновленная матрица Z.

Тип результата:

ArrayLike

Вспомогательные утилиты

Функции для загрузки/сохранения, работы с бэкендами (CPU/GPU) и генерации масок.

class denoise_bot.ML_component.utils.ArrayLike

Псевдоним для массивов, которые могут быть на CPU (NumPy) или GPU (CuPy).

псевдоним для ndarray | ndarray

class denoise_bot.ML_component.utils.BackendModule

Псевдоним для вычислительного бэкенда (модуль numpy или cupy).

псевдоним для ModuleType

denoise_bot.ML_component.utils.as_backend_array(array: ArrayLike, backend: BackendModule) ArrayLike[исходный код]

Универсальная функция для преобразования массива в массив нужного бэкенда.

Параметры:
  • array (ArrayLike) – Входной массив.

  • backend (BackendModule) – Целевой бэкенд (модуль numpy или cupy).

Результат:

Массив, находящийся на целевом устройстве (CPU или GPU).

Тип результата:

ArrayLike

denoise_bot.ML_component.utils.as_numpy(array: ArrayLike) ndarray[исходный код]

Универсальная функция для преобразования CuPy/NumPy массива в NumPy массив (на CPU).

Параметры:

array (ArrayLike) – Входной массив, который может быть numpy.ndarray или cupy.ndarray.

Результат:

Массив в формате NumPy.

Тип результата:

np.ndarray

denoise_bot.ML_component.utils.generate_mask(shape: Tuple[int, int], known_pixel_ratio: float, seed: int = 42) ndarray[исходный код]

Генерирует бинарную NumPy маску, где True означает известный пиксель.

Параметры:
  • shape (Tuple[int, int]) – Форма матрицы/изображения (высота, ширина).

  • known_pixel_ratio (float) – Доля известных пикселей (от 0 до 1).

  • seed (int) – Зерно для генератора случайных чисел для воспроизводимости.

Результат:

Булева NumPy маска той же формы, что и shape.

Тип результата:

np.ndarray

Исключение:

ValueError – Если known_pixel_ratio находится вне диапазона [0, 1].

denoise_bot.ML_component.utils.get_backend(use_gpu: bool = True) BackendModule[исходный код]

Выбор бэкенда.

Возвращает вычислительный бэкенд (NumPy или CuPy) в зависимости от доступности GPU и выбора пользователя.

Параметры:

use_gpu (bool) – Флаг, указывающий, нужно ли пытаться использовать GPU.

Результат:

Модуль numpy или cupy, который будет использоваться для вычислений.

Тип результата:

BackendModule

denoise_bot.ML_component.utils.load_image(source: str | IO[bytes], normalize: bool = True) ndarray[исходный код]

Загружает изображение из файла или байтового потока.

Параметры:
  • source (Union[str, IO[bytes]]) – Путь к файлу (str) или байтовый поток (io.BytesIO).

  • normalize (bool) – Если True, нормализует значения пикселей в диапазон [0, 1].

Результат:

Загруженное изображение в виде NumPy массива (на CPU).

Тип результата:

np.ndarray

Исключение:

ValueError – Если не удалось декодировать изображение из предоставленного источника.

denoise_bot.ML_component.utils.save_image(image_array: ArrayLike, path: str)[исходный код]

Сохраняет изображение из NumPy или CuPy массива в файл.

Параметры:
  • image_array (ArrayLike) – Массив изображения для сохранения.

  • path (str) – Путь для сохранения файла.

Функции для визуализации

Создание сравнительных графиков и графиков сходимости.

denoise_bot.ML_component.visualize.save_convergence_plot(norm_histories: List[List[float]], titles: List[str], colors: List[str], output_path: str)[исходный код]

Сохраняет графики сходимости ядерной нормы для каждого канала в файл.

Параметры:
  • norm_histories (List[List[float]]) – Список историй сходимости (каждая история - список чисел).

  • titles (List[str]) – Список заголовков для каждого графика.

  • colors (List[str]) – Список цветов для каждого графика.

  • output_path (str) – Путь для сохранения файла.

denoise_bot.ML_component.visualize.save_results_comparison(original: ArrayLike, damaged: ArrayLike, recovered: ArrayLike, output_path: str)[исходный код]

Сохранение сравнение результатов.

Сохраняет сравнение оригинального, поврежденного и восстановленного изображений в один файл.

Параметры:
  • original (ArrayLike) – Оригинальное изображение.

  • damaged (ArrayLike) – Изображение с зануленными пикселями.

  • recovered (ArrayLike) – Изображение после восстановления.

  • output_path (str) – Путь для сохранения файла.

denoise_bot.ML_component.visualize.show_image(ax: Axes, image: ArrayLike, title: str = '', grid: bool = False)[исходный код]

Отображение заданной картинки.

Отображает одно изображение на предоставленной оси matplotlib.Axes. Автоматически конвертирует CuPy массив в NumPy, если необходимо.

Параметры:
  • ax (plt.Axes) – Ось Matplotlib, на которой будет отрисовано изображение.

  • image (ArrayLike) – Массив изображения (numpy или cupy).

  • title (str) – Заголовок для изображения.

  • grid (bool) – Показывать ли сетку.