import gradio as gr
import torch
import numpy as np
import cv2
from PIL import Image
from spandrel import ModelLoader
import requests
import os

# 1. Загрузка модели в стиле "сохранение на диск"
MODEL_PATH = "realesrgan_x2.pth"

def load_model():
    if not os.path.exists(MODEL_PATH):
        d = "github.com"
        p = "xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"
        url = f"https://{d}/{p}"
        response = requests.get(url)
        with open(MODEL_PATH, "wb") as f:
            f.write(response.content)
    
    loader = ModelLoader()
    model_instance = loader.load_from_file(MODEL_PATH)
    model_instance.eval()
    return model_instance

model = load_model()

def recover(img, h_val, gamma_val, sharp_val):
    if img is None: return None
    
    img_np = np.array(img.convert("RGB"))
    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
    
    # 1. ЧИСТКА: используем ползунок h_val
    img_clean = cv2.fastNlMeansDenoisingColored(img_bgr, None, h_val, h_val, 7, 21)
    
    # ИИ часть
    img_rgb_ai = cv2.cvtColor(img_clean, cv2.COLOR_BGR2RGB)
    input_tensor = torch.from_numpy(img_rgb_ai).permute(2, 0, 1).float().divide(255).unsqueeze(0)
    with torch.no_grad():
        output_tensor = model(input_tensor)
    
    out_np = output_tensor.squeeze(0).permute(1, 2, 0).clamp(0, 1).numpy()
    out_np = (out_np * 255).astype(np.uint8)
    res_bgr = cv2.cvtColor(out_np, cv2.COLOR_RGB2BGR)
    
    # 2. ЗАСВЕТКА: используем ползунок gamma_val
    invGamma = 1.0 / gamma_val
    table = np.array([((i / 255.0) ** invGamma) * 255 for i in np.arange(0, 256)]).astype("uint8")
    res_bgr = cv2.LUT(res_bgr, table)
    
    # 3. РЕЗКОСТЬ: используем ползунок sharp_val
    blurred = cv2.GaussianBlur(res_bgr, (0, 0), 3)
    # Вычисляем коэффициенты на основе слайдера
    final_bgr = cv2.addWeighted(res_bgr, sharp_val, blurred, 1.0 - sharp_val, 0)
    final_bgr = np.clip(final_bgr, 0, 255).astype(np.uint8)
    
    # Финал
    final_rgb = cv2.cvtColor(final_bgr, cv2.COLOR_BGR2RGB)
    result_image = Image.fromarray(final_rgb)
    
    output_path = "result.png"
    result_image.save(output_path, "PNG")
    
    return result_image, output_path

# НАСТРОЙКА ИНТЕРФЕЙСА С ПОЛЗУНКАМИ
demo = gr.Interface(
    fn=recover, 
    inputs=[
        gr.Image(type="pil", label="Загрузить фото"),
        gr.Slider(minimum=1, maximum=15, value=5, step=1, label="Сила чистки (убирает зерно)"),
        gr.Slider(minimum=0.5, maximum=2.0, value=1.2, step=0.1, label="Гамма (против засветки)"),
        gr.Slider(minimum=1.0, maximum=2.5, value=1.4, step=0.1, label="Резкость линий")
    ], 
    outputs=[
        gr.Image(type="pil", label="Результат"), 
        gr.File(label="Скачать PNG")
    ],
    title="Artifact Remover PRO: Ручная настройка",
    description="Подкручивай ползунки, чтобы найти идеальный баланс между текстурой и четкостью."
)

if __name__ == "__main__":
    demo.launch()