Spaces:
Sleeping
Sleeping
| 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() | |