Spaces:
Runtime error
Runtime error
| 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() |