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

# Load the model into memory to make running multiple predictions efficient
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)