Spaces:
Sleeping
Sleeping
File size: 2,744 Bytes
40617c1 a840e7e 212761e b62e3d5 5920685 40617c1 5920685 40617c1 a840e7e 5920685 40617c1 5920685 212761e 40617c1 a840e7e 212761e a840e7e 212761e 5920685 212761e 40617c1 5920685 212761e a840e7e 212761e a840e7e 40617c1 212761e 3a27096 3d48cb6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | 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()
|