KirillDubov/nafnet-underwater

Underwater NAFNet (UIEB)

Модель для восстановления качества подводных фотографий, дообученная на датасете UIEB (Underwater Image Enhancement Benchmark).

Описание

Данная модель использует архитектуру NAFNet (Nonlinear Activation Free Network), которая показывает SOTA-результаты в задачах восстановления изображений. Модель эффективно корректирует цветовые искажения (потерю красного спектра) и удаляет «муть», характерную для съемок под водой.

Технические характеристики

• Архитектура: NAFNet (width=32)
• Датасет: UIEB (890 пар изображений)
• Предназначение: Улучшение подводных фото, сделанных на любительские камеры фото и экшн камеры.

Как использовать

Для загрузки модели из этого репозитория в вашем Python-скрипте:

from huggingface_hub import hf_hub_download
import torch

# Загрузка весов
weights_path = hf_hub_download(repo_id="KirillDubov/nafnet-underwater", filename="nafnet_uieb.pth")

# Дальше загружаете веса в свою архитектуру NAFNet:
# model.load_state_dict(torch.load(weights_path, map_location='cpu')['model'])