|
|
import os |
|
|
import cv2 |
|
|
import tempfile |
|
|
from modelscope.outputs import OutputKeys |
|
|
from modelscope.pipelines import pipeline |
|
|
from modelscope.utils.constant import Tasks |
|
|
from pathlib import Path |
|
|
import gradio as gr |
|
|
import numpy as np |
|
|
from PIL import Image, ImageEnhance |
|
|
|
|
|
|
|
|
img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization') |
|
|
|
|
|
def colorize_image(img_path): |
|
|
image = cv2.imread(str(img_path)) |
|
|
output = img_colorization(image[..., ::-1]) |
|
|
result = output[OutputKeys.OUTPUT_IMG].astype(np.uint8) |
|
|
temp_dir = tempfile.mkdtemp() |
|
|
out_path = os.path.join(temp_dir, 'colorized.png') |
|
|
cv2.imwrite(out_path, result) |
|
|
return out_path |
|
|
|
|
|
def enhance_image(img_path, brightness=1.0, contrast=1.0): |
|
|
image = Image.open(img_path) |
|
|
enhancer_brightness = ImageEnhance.Brightness(image) |
|
|
image = enhancer_brightness.enhance(brightness) |
|
|
enhancer_contrast = ImageEnhance.Contrast(image) |
|
|
image = enhancer_contrast.enhance(contrast) |
|
|
temp_dir = tempfile.mkdtemp() |
|
|
enhanced_path = os.path.join(temp_dir, 'enhanced.png') |
|
|
image.save(enhanced_path) |
|
|
return enhanced_path |
|
|
|
|
|
def process_image(img_path, brightness, contrast): |
|
|
colorized_path = colorize_image(img_path) |
|
|
enhanced_path = enhance_image(colorized_path, brightness, contrast) |
|
|
return [img_path, enhanced_path], enhanced_path |
|
|
|
|
|
title = "🌈 Color Restorization Model" |
|
|
description = "Upload a black & white photo to restore it in color using a deep learning model." |
|
|
|
|
|
with gr.Blocks(title=title) as demo: |
|
|
gr.Markdown(f"## {title}") |
|
|
gr.Markdown(description) |
|
|
|
|
|
with gr.Row(): |
|
|
with gr.Column(): |
|
|
input_image = gr.Image(type="filepath", label="Upload B&W Image") |
|
|
brightness_slider = gr.Slider(0.5, 2.0, value=1.0, label="Brightness") |
|
|
contrast_slider = gr.Slider(0.5, 2.0, value=1.0, label="Contrast") |
|
|
submit_btn = gr.Button("Colorize") |
|
|
with gr.Column(): |
|
|
comparison = gr.Gallery(label="Original vs Colorized").style(grid=[2], height="auto") |
|
|
download_btn = gr.File(label="Download Colorized Image") |
|
|
|
|
|
submit_btn.click( |
|
|
fn=process_image, |
|
|
inputs=[input_image, brightness_slider, contrast_slider], |
|
|
outputs=[comparison, download_btn] |
|
|
) |
|
|
|
|
|
demo.launch(enable_queue=True) |
|
|
|
