UpscaleThis / app.py
Doubleupai's picture
Update app.py
5fa6790 verified
import gradio as gr
from PIL import Image
import torch
from torchvision import transforms
# Загрузка предобученной модели ESRGAN через torch.hub
def load_model():
model = torch.hub.load('pytorch/vision:v0.10.0', 'esrgan', pretrained=True)
model.eval()
return model
# Функция для увеличения разрешения изображения
def upscale_image(image):
# Преобразование изображения в тензор
input_image = transforms.ToTensor()(image).unsqueeze(0)
# Увеличение разрешения с помощью модели
with torch.no_grad():
output_image = model(input_image)
# Преобразование тензора обратно в изображение
output_image = output_image.squeeze(0).clamp(0, 1)
output_image = transforms.ToPILImage()(output_image)
return output_image
# Загрузка модели
model = load_model()
# Создание интерфейса Gradio
iface = gr.Interface(
fn=upscale_image,
inputs=gr.Image(type="pil", label="Входное изображение"),
outputs=gr.Image(type="pil", label="Увеличенное изображение"),
title="AI Image Upscaler (ESRGAN)",
description="Загрузите изображение, чтобы увеличить его разрешение с помощью ESRGAN."
)
# Запуск интерфейса
iface.launch()