blend_processor.py

import gradio as gr
import numpy as np
import cv2
from PIL import Image, ImageChops
from skimage import img_as_float, img_as_ubyte
import copy
from typing import Optional, Union

def cv22ski(cv2_image: np.ndarray) -> np.ndarray:
    """Convert CV2 image to skimage float format"""
    return img_as_float(cv2_image)

def ski2cv2(ski: np.ndarray) -> np.ndarray:
    """Convert skimage float format to CV2 image"""
    return img_as_ubyte(ski)

def cv22pil(cv2_img: np.ndarray) -> Image.Image:
    """Convert CV2 image to PIL Image"""
    cv2_img = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB)
    return Image.fromarray(cv2_img)

def pil2cv2(pil_img: Image.Image) -> np.ndarray:
    """Convert PIL Image to CV2 image"""
    np_img_array = np.asarray(pil_img)
    return cv2.cvtColor(np_img_array, cv2.COLOR_RGB2BGR)

def blend_color_burn(background_image: Image.Image, layer_image: Image.Image) -> Image.Image:
    """Apply color burn blend mode"""
    img_1 = cv22ski(pil2cv2(background_image))
    img_2 = cv22ski(pil2cv2(layer_image))
    img = 1 - (1 - img_2) / (img_1 + 0.001)
    mask_1 = img < 0
    mask_2 = img > 1
    img = img * (1 - mask_1)
    img = img * (1 - mask_2) + mask_2
    return cv22pil(ski2cv2(img))

def blend_color_dodge(background_image: Image.Image, layer_image: Image.Image) -> Image.Image:
    """Apply color dodge blend mode"""
    img_1 = cv22ski(pil2cv2(background_image))
    img_2 = cv22ski(pil2cv2(layer_image))
    img = img_2 / (1.0 - img_1 + 0.001)
    mask_2 = img > 1
    img = img * (1 - mask_2) + mask_2
    return cv22pil(ski2cv2(img))

def blend_linear_burn(background_image: Image.Image, layer_image: Image.Image) -> Image.Image:
    """Apply linear burn blend mode"""
    img_1 = cv22ski(pil2cv2(background_image))
    img_2 = cv22ski(pil2cv2(layer_image))
    img = img_1 + img_2 - 1
    mask_1 = img < 0
    img = img * (1 - mask_1)
    return cv22pil(ski2cv2(img))

def blend_linear_dodge(background_image: Image.Image, layer_image: Image.Image) -> Image.Image:
    """Apply linear dodge blend mode"""
    img_1 = cv22ski(pil2cv2(background_image))
    img_2 = cv22ski(pil2cv2(layer_image))
    img = img_1 + img_2
    mask_2 = img > 1
    img = img * (1 - mask_2) + mask_2
    return cv22pil(ski2cv2(img))

def blend_lighten(background_image: Image.Image, layer_image: Image.Image) -> Image.Image:
    """Apply lighten blend mode"""
    img_1 = cv22ski(pil2cv2(background_image))
    img_2 = cv22ski(pil2cv2(layer_image))
    img = img_1 - img_2
    mask = img > 0
    img = img_1 * mask + img_2 * (1 - mask)
    return cv22pil(ski2cv2(img))

def blend_dark(background_image: Image.Image, layer_image: Image.Image) -> Image.Image:
    """Apply darken blend mode"""
    img_1 = cv22ski(pil2cv2(background_image))
    img_2 = cv22ski(pil2cv2(layer_image))
    img = img_1 - img_2
    mask = img < 0
    img = img_1 * mask + img_2 * (1 - mask)
    return cv22pil(ski2cv2(img))

def blend_screen(background_image: Image.Image, layer_image: Image.Image) -> Image.Image:
    """Apply screen blend mode"""
    img_1 = cv22ski(pil2cv2(background_image))
    img_2 = cv22ski(pil2cv2(layer_image))
    img = 1 - (1 - img_1) * (1 - img_2)
    return cv22pil(ski2cv2(img))

def blend_overlay(background_image: Image.Image, layer_image: Image.Image) -> Image.Image:
    """Apply overlay blend mode"""
    img_1 = cv22ski(pil2cv2(background_image))
    img_2 = cv22ski(pil2cv2(layer_image))
    mask = img_2 < 0.5
    img = 2 * img_1 * img_2 * mask + (1 - mask) * (1 - 2 * (1 - img_1) * (1 - img_2))
    return cv22pil(ski2cv2(img))

def blend_soft_light(background_image: Image.Image, layer_image: Image.Image) -> Image.Image:
    """Apply soft light blend mode"""
    img_1 = cv22ski(pil2cv2(background_image))
    img_2 = cv22ski(pil2cv2(layer_image))
    mask = img_1 < 0.5
    T1 = (2 * img_1 - 1) * (img_2 - img_2 * img_2) + img_2
    T2 = (2 * img_1 - 1) * (np.sqrt(img_2) - img_2) + img_2
    img = T1 * mask + T2 * (1 - mask)
    return cv22pil(ski2cv2(img))

def blend_hard_light(background_image: Image.Image, layer_image: Image.Image) -> Image.Image:
    """Apply hard light blend mode"""
    img_1 = cv22ski(pil2cv2(background_image))
    img_2 = cv22ski(pil2cv2(layer_image))
    mask = img_1 < 0.5
    T1 = 2 * img_1 * img_2
    T2 = 1 - 2 * (1 - img_1) * (1 - img_2)
    img = T1 * mask + T2 * (1 - mask)
    return cv22pil(ski2cv2(img))

def blend_vivid_light(background_image: Image.Image, layer_image: Image.Image) -> Image.Image:
    """Apply vivid light blend mode"""
    img_1 = cv22ski(pil2cv2(background_image))
    img_2 = cv22ski(pil2cv2(layer_image))
    mask = img_1 < 0.5
    T1 = 1 - (1 - img_2) / (2 * img_1 + 0.001)
    T2 = img_2 / (2 * (1 - img_1) + 0.001)
    mask_1 = T1 < 0
    mask_2 = T2 > 1
    T1 = T1 * (1 - mask_1)
    T2 = T2 * (1 - mask_2) + mask_2
    img = T1 * mask + T2 * (1 - mask)
    return cv22pil(ski2cv2(img))

def blend_pin_light(background_image: Image.Image, layer_image: Image.Image) -> Image.Image:
    """Apply pin light blend mode"""
    img_1 = cv22ski(pil2cv2(background_image))
    img_2 = cv22ski(pil2cv2(layer_image))
    mask_1 = img_2 < (img_1 * 2 - 1)
    mask_2 = img_2 > 2 * img_1
    T1 = 2 * img_1 - 1
    T2 = img_2
    T3 = 2 * img_1
    img = T1 * mask_1 + T2 * (1 - mask_1) * (1 - mask_2) + T3 * mask_2
    return cv22pil(ski2cv2(img))

def blend_linear_light(background_image: Image.Image, layer_image: Image.Image) -> Image.Image:
    """Apply linear light blend mode"""
    img_1 = cv22ski(pil2cv2(background_image))
    img_2 = cv22ski(pil2cv2(layer_image))
    img = img_2 + img_1 * 2 - 1
    mask_1 = img < 0
    mask_2 = img > 1
    img = img * (1 - mask_1)
    img = img * (1 - mask_2) + mask_2
    return cv22pil(ski2cv2(img))

def blend_hard_mix(background_image: Image.Image, layer_image: Image.Image) -> Image.Image:
    """Apply hard mix blend mode"""
    img_1 = cv22ski(pil2cv2(background_image))
    img_2 = cv22ski(pil2cv2(layer_image))
    img = img_1 + img_2
    mask = img_1 + img_2 > 1
    img = img * (1 - mask) + mask
    img = img * mask
    return cv22pil(ski2cv2(img))

def chop_image(background_image: Image.Image, layer_image: Image.Image, blend_mode: str, opacity: int) -> Image.Image:
    """Apply blend mode and opacity to images"""
    ret_image = background_image
    
    blend_functions = {
        'normal': lambda: copy.deepcopy(layer_image),
        'multiply': lambda: ImageChops.multiply(background_image, layer_image),
        'screen': lambda: ImageChops.screen(background_image, layer_image),
        'add': lambda: ImageChops.add(background_image, layer_image, 1, 0),
        'subtract': lambda: ImageChops.subtract(background_image, layer_image, 1, 0),
        'difference': lambda: ImageChops.difference(background_image, layer_image),
        'darker': lambda: ImageChops.darker(background_image, layer_image),
        'lighter': lambda: ImageChops.lighter(background_image, layer_image),
        'color_burn': lambda: blend_color_burn(background_image, layer_image),
        'color_dodge': lambda: blend_color_dodge(background_image, layer_image),
        'linear_burn': lambda: blend_linear_burn(background_image, layer_image),
        'linear_dodge': lambda: blend_linear_dodge(background_image, layer_image),
        'overlay': lambda: blend_overlay(background_image, layer_image),
        'soft_light': lambda: blend_soft_light(background_image, layer_image),
        'hard_light': lambda: blend_hard_light(background_image, layer_image),
        'vivid_light': lambda: blend_vivid_light(background_image, layer_image),
        'pin_light': lambda: blend_pin_light(background_image, layer_image),
        'linear_light': lambda: blend_linear_light(background_image, layer_image),
        'hard_mix': lambda: blend_hard_mix(background_image, layer_image)
    }
    
    if blend_mode in blend_functions:
        ret_image = blend_functions[blend_mode]()
    
    # Apply opacity
    if opacity == 0:
        ret_image = background_image
    elif opacity < 100:
        alpha = 1.0 - float(opacity) / 100
        ret_image = Image.blend(ret_image, background_image, alpha)
    
    return ret_image

def process_images(background: Optional[np.ndarray], 
                  layer: Optional[np.ndarray], 
                  blend_mode: str,
                  opacity: float) -> Optional[np.ndarray]:
    """Process images with selected blend mode and opacity"""
    if background is None or layer is None:
        return None
        
    # Convert numpy arrays to PIL Images
    background_pil = Image.fromarray(background)
    layer_pil = Image.fromarray(layer)
    
    # Ensure images are in RGB mode
    background_pil = background_pil.convert('RGB')
    layer_pil = layer_pil.convert('RGB')
    
    # Apply blend mode
    result = chop_image(background_pil, layer_pil, blend_mode, int(opacity * 100))
    
    # Convert back to numpy array
    return np.array(result)

def create_blend_tab():
    """Create the blend modes tab interface"""
    with gr.Tab("Blend Modes"):
        with gr.Row():
            with gr.Column():
                background_image = gr.Image(label="Background Image", height=256)
                layer_image = gr.Image(label="Layer Image", height=256)
                
                blend_mode = gr.Dropdown(
                    choices=[
                        "normal", "multiply", "screen", "overlay",
                        "soft_light", "hard_light", "color_burn", "color_dodge",
                        "linear_burn", "linear_dodge", "vivid_light", "linear_light",
                        "pin_light", "hard_mix", "difference", "add", "subtract",
                        "darker", "lighter"
                    ],
                    value="normal",
                    label="Blend Mode"
                )
                
                opacity = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=1.0,
                    step=0.01,
                    label="Opacity"
                )
                
                blend_btn = gr.Button("Apply Blend")
                
            with gr.Column():
                output_image = gr.Image(label="Blended Image")
        
        blend_btn.click(
            fn=process_images,
            inputs=[background_image, layer_image, blend_mode, opacity],
            outputs=output_image
        )