import os
import cv2
import numpy as np
import gradio as gr

# Load pretrained models
proto_path = "colorization_deploy_v2.prototxt"
model_path = "colorization_release_v2.caffemodel"
pts_path = "pts_in_hull.npy"

net = cv2.dnn.readNetFromCaffe(proto_path, model_path)
pts = np.load(pts_path)
pts = pts.transpose().reshape(2, 313, 1, 1).astype(np.float32)
net.getLayer(net.getLayerId("class8_ab")).blobs = [pts]
net.getLayer(net.getLayerId("conv8_313_rh")).blobs = [np.full((1, 313), 2.606, np.float32)]

def colorize(img_bgr):
    h, w = img_bgr.shape[:2]
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    img_lab = cv2.cvtColor(img_rgb.astype("float32") / 255.0, cv2.COLOR_RGB2LAB)
    L = img_lab[:, :, 0]
    L_rs = cv2.resize(L, (224, 224))
    L_rs -= 50
    net.setInput(cv2.dnn.blobFromImage(L_rs))
    ab_dec = net.forward()[0].transpose((1, 2, 0))
    ab_dec_us = cv2.resize(ab_dec, (w, h))
    lab_output = np.concatenate((L[:, :, np.newaxis], ab_dec_us), axis=2)
    bgr_output = cv2.cvtColor(lab_output.astype("uint8"), cv2.COLOR_LAB2BGR)
    return bgr_output

def restore(img, enhance_face, remove_scratches, apply_colorization):
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    if remove_scratches:
        img = cv2.medianBlur(img, 3)
    if enhance_face:
        blur = cv2.GaussianBlur(img, (0, 0), 3)
        img = cv2.addWeighted(img, 1.5, blur, -0.5, 0)
    if apply_colorization:
        img = colorize(img)
    restored = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    cv2.imwrite("restored_output.png", restored)
    return restored, "restored_output.png"

with gr.Blocks(title="AI Old Photo Restorer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🧓 AI Old Photo Restorer")
    gr.Markdown("Give new life to your old, damaged, or black-and-white photos with facial enhancement, scratch removal, and colorization!")

    with gr.Row():
        with gr.Column(scale=1):
            inp = gr.Image(label="📤 Upload Old Photo", type="numpy", image_mode='RGB')
            enhance_face = gr.Checkbox(label="✨ Enhance Facial Details", value=True)
            remove_scratches = gr.Checkbox(label="🩹 Remove Scratches", value=True)
            apply_colorization = gr.Checkbox(label="🎨 Colorize Black & White Photo", value=True)
            btn = gr.Button("🔁 Restore Photo", variant="primary")
            clear_btn = gr.ClearButton([inp, enhance_face, remove_scratches, apply_colorization])
        with gr.Column(scale=1):
            out_image = gr.Image(label="✅ Restored Output", interactive=False)
            download = gr.File(label="⬇️ Download Restored Image")

    btn.click(restore, [inp, enhance_face, remove_scratches, apply_colorization], [out_image, download])

demo.launch()