File size: 4,706 Bytes
749f522 |
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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 |
import os
import cv2
import tempfile
from pathlib import Path
import gradio as gr
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
# Load the colorization model into memory once at startup
img_colorization = pipeline(
Tasks.image_colorization,
model="iic/cv_ddcolor_image-colorization"
)
def colorize_image(img_path: str) -> str:
"""
Reads a B&W image from disk, runs the colorization model,
writes the colorized result to a temp file, and returns its 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: str,
brightness: float = 1.0,
contrast: float = 1.0,
edge_enhance: bool = False
) -> str:
"""
Opens a colorized image from disk, applies brightness, contrast,
and optional edge enhancement, saves to a temp file, and returns its path.
"""
image = Image.open(img_path)
# Adjust brightness
image = ImageEnhance.Brightness(image).enhance(brightness)
# Adjust contrast
image = ImageEnhance.Contrast(image).enhance(contrast)
# Optionally apply an edge enhancement filter
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: str,
brightness: float,
contrast: float,
edge_enhance: bool,
output_format: str
):
"""
1) Colorizes the uploaded B&W image.
2) Applies the chosen brightness/contrast/edge-enhancement.
3) Re‐saves in the user’s chosen format (PNG/JPEG/TIFF).
Returns:
- A list [original_path, final_path] for side-by-side display.
- The final image’s file path for download.
"""
# Step 1: colorize
colorized_path = colorize_image(img_path)
# Step 2: enhancement
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 ([original, enhanced], download_path)
return ([img_path, enhanced_path], output_path)
# Title and description shown at the top of the interface
TITLE = "🌈 Color Restorization Model"
DESCRIPTION = "Upload a black & white photo to restore it in color using a deep learning model."
# Build the Gradio Blocks interface
with gr.Blocks(title=TITLE) as app:
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" # Enables zoom/pan on the uploaded image
)
brightness_slider = gr.Slider(
minimum=0.5, maximum=2.0, value=1.0,
label="Brightness"
)
contrast_slider = gr.Slider(
minimum=0.5, maximum=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, # two images side by side
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]
)
# Launch in “production” mode: bind to 0.0.0.0 on configurable port
if __name__ == "__main__":
port = int(os.environ.get("PORT", 7860))
app.queue().launch(server_name="0.0.0.0", server_port=port)
|