File size: 3,341 Bytes
90bcc4d 24f487e 90bcc4d 24f487e 90bcc4d 24f487e 90bcc4d 24f487e 90bcc4d 24f487e 90bcc4d 24f487e 90bcc4d 24f487e 90bcc4d 24f487e 90bcc4d 24f487e 90bcc4d |
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 |
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, ImageFilter
# Load the model once at startup for efficiency
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, edge_enhance=False):
image = Image.open(img_path)
# Adjust brightness
image = ImageEnhance.Brightness(image).enhance(brightness)
# Adjust contrast
image = ImageEnhance.Contrast(image).enhance(contrast)
# Optionally apply edge enhancement
if edge_enhance:
image = image.filter(ImageFilter.EDGE_ENHANCE)
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, edge_enhance, output_format):
# Step 1: colorize
colorized_path = colorize_image(img_path)
# Step 2: enhance (brightness / contrast / edge)
enhanced_path = enhance_image(colorized_path, brightness, contrast, edge_enhance)
# Step 3: convert to chosen format
img = Image.open(enhanced_path)
temp_dir = tempfile.mkdtemp()
filename = f'colorized_image.{output_format.lower()}'
output_path = os.path.join(temp_dir, filename)
img.save(output_path, format=output_format.upper())
# Return side-by-side (original, enhanced) plus the single downloadable file
return ([img_path, enhanced_path], output_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", tool="editor")
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")
edge_enhance_checkbox = gr.Checkbox(label="Apply Edge Enhancement")
output_format_dropdown = gr.Dropdown(
choices=["PNG", "JPEG", "TIFF"],
value="PNG",
label="Output Format"
)
submit_btn = gr.Button("Colorize")
with gr.Column():
comparison_gallery = gr.Gallery(
label="Original vs Colorized",
columns=2,
height="auto"
)
download_btn = gr.File(label="Download Colorized Image")
submit_btn.click(
fn=process_image,
inputs=[input_image, brightness_slider, contrast_slider, edge_enhance_checkbox, output_format_dropdown],
outputs=[comparison_gallery, download_btn]
)
demo.launch(enable_queue=True)
|