File size: 3,087 Bytes
c328ebc |
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 |
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 into memory to make running multiple predictions efficient
img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization')
def apply_clahe(image):
# Convert to LAB color space
lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)
l, a, b = cv2.split(lab)
# Apply CLAHE to the L-channel
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
cl = clahe.apply(l)
# Merge the CLAHE enhanced L-channel back with A and B channels
limg = cv2.merge((cl,a,b))
# Convert back to RGB color space
final = cv2.cvtColor(limg, cv2.COLOR_LAB2RGB)
return final
def enhance_image_pil(image_pil, brightness=1.0, contrast=1.0):
# Apply brightness and contrast enhancements
enhancer_brightness = ImageEnhance.Brightness(image_pil)
image_pil = enhancer_brightness.enhance(brightness)
enhancer_contrast = ImageEnhance.Contrast(image_pil)
image_pil = enhancer_contrast.enhance(contrast)
# Apply edge enhancement
image_pil = image_pil.filter(ImageFilter.EDGE_ENHANCE)
return image_pil
def process_image(img_path, brightness, contrast):
# Read the input image
image = cv2.imread(str(img_path))
# Colorize the image
output = img_colorization(image[..., ::-1])
result = output[OutputKeys.OUTPUT_IMG].astype(np.uint8)
# Apply CLAHE
result_clahe = apply_clahe(result)
# Convert to PIL Image for further enhancements
result_pil = Image.fromarray(result_clahe)
# Apply brightness, contrast, and edge enhancements
enhanced_pil = enhance_image_pil(result_pil, brightness, contrast)
# Save the enhanced image to a temporary file
temp_dir = tempfile.mkdtemp()
enhanced_path = os.path.join(temp_dir, 'enhanced.png')
enhanced_pil.save(enhanced_path)
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)
|