app.py

import argparse
import sys
import numpy as np
import os
import cv2
import skimage
import gradio as gr
from pygltflib import GLTF2
from pygltflib.utils import ImageFormat, Image
import PIL.Image
import io
import base64

METHODS = ('lhm', 'pccm', 'reinhard')


def transfer_lhm(content, reference):
    """Transfers colors from a reference image to a content image using the
    Linear Histogram Matching.

    content: NumPy array (HxWxC)
    reference: NumPy array (HxWxC)
    """
    # Convert HxWxC image to a (H*W)xC matrix.
    shape = content.shape
    assert len(shape) == 3
    content = content.reshape(-1, shape[-1]).astype(np.float32)
    reference = reference.reshape(-1, shape[-1]).astype(np.float32)

    def matrix_sqrt(X):
        eig_val, eig_vec = np.linalg.eig(X)
        return eig_vec.dot(np.diag(np.sqrt(eig_val))).dot(eig_vec.T)
    
    #
    mu_content = np.mean(content, axis=0)
    #
    mu_reference = np.mean(reference, axis=0)

    cov_content = np.cov(content, rowvar=False)
    cov_reference = np.cov(reference, rowvar=False)

    #
    result = matrix_sqrt(cov_reference)

    #
    result = result.dot(np.linalg.inv(matrix_sqrt(cov_content)))

    #
    result = result.dot((content - mu_content).T).T
    #result = result.dot((content*1 - mu_content*0.5).T).T*3

    #
    result = result + mu_reference


    # Restore image dimensions.
    result = result.reshape(shape).clip(0, 255).round().astype(np.uint8)

    return result


def transfer_pccm(content, reference):
    """Transfers colors from a reference image to a content image using
    Principal Component Color Matching.

    content: NumPy array (HxWxC)
    reference: NumPy array (HxWxC)
    """
    # Convert HxWxC image to a (H*W)xC matrix.
    shape = content.shape
    assert len(shape) == 3
    content = content.reshape(-1, shape[-1]).astype(np.float32)
    reference = reference.reshape(-1, shape[-1]).astype(np.float32)

    mu_content = np.mean(content, axis=0)
    mu_reference = np.mean(reference, axis=0)

    cov_content = np.cov(content, rowvar=False)
    cov_reference = np.cov(reference, rowvar=False)

    eigval_content, eigvec_content = np.linalg.eig(cov_content)
    eigval_reference, eigvec_reference = np.linalg.eig(cov_reference)

    scaling = np.diag(np.sqrt(eigval_reference / eigval_content))
    transform = eigvec_reference.dot(scaling).dot(eigvec_content.T)
    result = (content - mu_content).dot(transform.T) + mu_reference
    # Restore image dimensions.
    result = result.reshape(shape).clip(0, 255).round().astype(np.uint8)

    return result


def transfer_reinhard(content, reference):
    """Transfers colors from a reference image to a content image using the
    technique from Reinhard et al.

    content: NumPy array (HxWxC)
    reference: NumPy array (HxWxC)
    """
    # Convert HxWxC image to a (H*W)xC matrix.
    shape = content.shape
    assert len(shape) == 3
    content = content.reshape(-1, shape[-1]).astype(np.float32)
    reference = reference.reshape(-1, shape[-1]).astype(np.float32)

    m1 = np.array([
        [0.3811, 0.1967, 0.0241],
        [0.5783, 0.7244, 0.1288],
        [0.0402, 0.0782, 0.8444],
    ])

    m2 = np.array([
        [0.5774, 0.4082, 0.7071],
        [0.5774, 0.4082, -0.7071],
        [0.5774, -0.8165, 0.0000],
    ])

    m3 = np.array([
        [0.5774, 0.5774, 0.5774],
        [0.4082, 0.4082, -0.8165],
        [0.7071, -0.7071, 0.0000],
    ])

    m4 = np.array([
        [4.4679, -1.2186, 0.0497],
        [-3.5873, 2.3809, -0.2439],
        [0.1193, -0.1624, 1.2045],
    ])

    # Avoid log of 0. Clipping is used instead of adding epsilon, to avoid
    # taking a log of a small number whose very low output distorts the results.
    # WARN: This differs from the Reinhard paper, where no adjustment is made.
    lab_content = np.log10(np.maximum(1.0, content.dot(m1))).dot(m2)
    lab_reference = np.log10(np.maximum(1.0, reference.dot(m1))).dot(m2)

    mu_content = lab_content.mean(axis=0)  # shape=3
    mu_reference = lab_reference.mean(axis=0)

    std_source = np.std(content, axis=0)
    std_target = np.std(reference, axis=0)
    #variable percentage for mu and std
    result = lab_content - mu_content
    result *= std_target
    result /= std_source
    result += mu_reference
    result = (10 ** result.dot(m3)).dot(m4)
    # Restore image dimensions.
    result = result.reshape(shape).clip(0, 255).round().astype(np.uint8)

    return result




# =================================================================================




def runModel(content, style):
    if os.path.exists("0.png"):
        os.remove("0.png")
    gltf = GLTF2().load(content)
    gltf.convert_images(ImageFormat.FILE,"/tmp/", override=True)
    gltf.images[0].uri
    gltf_reinhard = GLTF2().load(content)
    gltf_lhm = GLTF2().load(content)
    gltf_pccm = GLTF2().load(content)
    ori_image=PIL.Image.open('/tmp/0.png')
    ori_image=ori_image.convert('RGB')
    content_img = np.array(ori_image)[:, :, :3]
    style_img = np.array(style)[:, :, :3]
    output_reinhard = transfer_reinhard(content_img, style_img)
    output_lhm = transfer_lhm(content_img, style_img)
    output_pccm = transfer_pccm(content_img, style_img)
    image_reinhard = PIL.Image.fromarray(output_reinhard,'RGB')
    image_lhm = PIL.Image.fromarray(output_lhm,'RGB')
    image_pccm = PIL.Image.fromarray(output_pccm,'RGB')
    image_reinhard=image_reinhard.save('/tmp/reinhard.png')
    image_lhm=image_lhm.save('/tmp/lhm.png')
    image_pccm=image_pccm.save('/tmp/pccm.png')
    image1=Image()
    image2=Image()
    image3=Image()
    image1.uri='/tmp/reinhard.png'
    image2.uri='/tmp/lhm.png'
    image3.uri='/tmp/pccm.png'
    gltf_reinhard.images[0]=image1
    gltf_lhm.images[0]=image2
    gltf_pccm.images[0]=image3
    gltf_reinhard.convert_images(ImageFormat.DATAURI)
    gltf_lhm.convert_images(ImageFormat.DATAURI)
    gltf_pccm.convert_images(ImageFormat.DATAURI)
    gltf_reinhard.images[0].uri 
    gltf_reinhard.images[0].name
    gltf_lhm.images[0].uri 
    gltf_lhm.images[0].name 
    gltf_pccm.images[0].uri 
    gltf_pccm.images[0].name
    gltf_pccm.save('/tmp/pccm.glb')
    gltf_lhm.save('/tmp/lhm.glb')
    gltf_reinhard.save('/tmp/reinhard.glb')
    return '/tmp/pccm.glb','/tmp/lhm.glb', '/tmp/reinhard.glb'

demo = gr.Interface(runModel, [gr.Model3D(),gr.Image(type ='pil')], outputs = [gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="Reinhard"), gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="lhm"), gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="pccm")], )

demo.launch()