File size: 2,790 Bytes
ca7e923 ed0b9b7 ca7e923 ed0b9b7 ca7e923 ed0b9b7 |
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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
import os
import cv2
import tempfile
import numpy as np
import gradio as gr
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from pathlib import Path
# Initialize model
def load_model():
global img_colorization
img_colorization = pipeline(
Tasks.image_colorization,
model='iic/cv_ddcolor_image-colorization',
model_revision='v1.0.0'
)
def inference(img):
if img is None:
raise gr.Error("Please upload an image first")
with tempfile.TemporaryDirectory() as temp_dir:
# Convert PIL image to numpy array if needed
if isinstance(img, np.ndarray):
image = img
else:
image = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
# Process image
output = img_colorization(image[..., ::-1])
result = output['output_img'].astype(np.uint8)
# Save result
out_path = os.path.join(temp_dir, 'colorized.png')
cv2.imwrite(out_path, result)
return Path(out_path), "✅ Colorization completed successfully!"
# Create modern UI with Blocks
with gr.Blocks(theme="soft", title="🎨 AI Color Restoration Studio") as demo:
gr.Markdown("""
# 🎨 AI Color Restoration Studio
Transform your black & white photos into vibrant colorized versions using state-of-the-art AI!
Upload an image and watch as our deep learning model automatically adds natural colors.
""")
with gr.Row():
with gr.Column(scale=1):
input_img = gr.Image(
label="Upload Monochrome Image",
type="pil",
height=400,
sources=["upload"],
interactive=True
)
submit_btn = gr.Button("✨ Colorize Image", variant="primary")
clear_btn = gr.ClearButton()
with gr.Column(scale=1):
output_img = gr.Image(
label="Colorized Result",
type="pil",
height=400,
interactive=False
)
download_btn = gr.File(label="Download Result")
status = gr.Textbox(label="Status", interactive=False)
# Examples section
gr.Examples(
examples=[
["examples/1.jpg"],
["examples/2.jpg"],
["examples/3.jpg"]
],
inputs=[input_img],
outputs=[output_img, status],
fn=inference,
cache_examples=True
)
# Event handlers
submit_btn.click(
fn=inference,
inputs=[input_img],
outputs=[output_img, status]
)
clear_btn.add([input_img, output_img, status])
if __name__ == "__main__":
load_model()
demo.launch(debug=True)
|