Spaces:
Running
Running
| import os | |
| os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" | |
| import gradio as gr | |
| import onnxruntime as ort | |
| import importlib.util | |
| from huggingface_hub import hf_hub_download | |
| import numpy as np | |
| import cv2 | |
| import psutil | |
| import time | |
| import sys | |
| token_ = os.environ.get("HF_TOKEN") | |
| model_repo_id = "heyh97791/NCF" | |
| model_filename = "ncf.onnx" | |
| file_path = "skin_protection.py" | |
| pretrained_name = "skin.tflite" | |
| protector_instance = None | |
| I18N = { | |
| "lang_btn": { | |
| "en": "中文", | |
| "zh": "English" | |
| }, | |
| "title": { | |
| "en": "# 🎨 ColorFM-L [CPU] ", | |
| "zh": "# 🎨 ColorFM-L [CPU] " | |
| }, | |
| "paper_code": { | |
| "en": """ | |
| 📄 **Paper:** [ColorFM: An Optimization-to-Learning Framework for Color Transfer via Flow Matching (ECCV 2026)](https://github.com/cszn/ColorFM) | 💻 **Code:** [GitHub](https://github.com/cszn/ColorFM) | ⚡ **Model:** Feed-forward model | |
| """, | |
| "zh": """ | |
| 📄 **论文：** [ColorFM: An Optimization-to-Learning Framework for Color Transfer via Flow Matching (ECCV 2026)](https://github.com/cszn/ColorFM) | 💻 **代码：** [GitHub](https://github.com/cszn/ColorFM) | ⚡ **模型：** 前馈式模型 | |
| """ | |
| }, | |
| "desc_header": { | |
| "en": "💡 What is Color Transfer?", | |
| "zh": "💡 追色是什么" | |
| }, | |
| "desc_content": { | |
| "en": """ | |
| **1. What is "Color Transfer"?** | |
| It is a technique that extracts the **color palette and atmosphere** from a reference image (Style) and applies it to your target image (Content), while preserving the original structure. | |
| **2. Use Cases:** | |
| * 📸 **Photography:** Instantly mimic color grading styles from master photographers. | |
| * 🎨 **Art & Design:** Unify the color theme of different assets quickly. | |
| **3. Tips for Best Results:** | |
| * ✅ **Match Content:** Results are best when the content and style images share similar scenes (e.g., Landscape to Landscape). | |
| * 🎨 **Hue Similarity:** Images with closer hues between content and style may produce better color transfer results. | |
| * ⚡ **High Resolution:** Supports processing high-resolution images (use the Settings menu to optimize speed). | |
| **🛡️ Privacy Disclaimer:** | |
| This demo runs entirely on the cloud instance. Your uploaded images are processed in memory and are **NOT saved** or stored permanently on our servers. | |
| """, | |
| "zh": """ | |
| **1. 什么是“追色”？** | |
| 追色（Color Transfer）是指从一张参考图（风格图）中提取**色调与氛围**，并将其“迁移”到你的目标图片（内容图）上，同时保留原图的细节纹理。 | |
| **2. 主要用途：** | |
| * 📸 **摄影后期：** 一键复刻摄影大师的调色风格，无需手动调参。 | |
| * 🎨 **设计创作：** 快速统一多张素材的色调风格，提高创作效率。 | |
| **3. 如何获得最佳效果？** | |
| * ✅ **内容匹配：** 当内容图与风格图的场景相似时（例如都是风景），效果通常最好。 | |
| * 🎨 **色相相似：** 内容图像和风格图像之间色相相似，可能会呈现更好的效果 | |
| * ⚡ **高清支持：** 支持高分辨率图片处理（可在“设置”中调整分辨率以获得更快速度）。 | |
| **🛡️ 免责声明：** | |
| 本Demo仅供演示体验。您的图片仅在内存中进行临时处理，**不会被保存**、存储或用于任何其他用途，处理结束后即刻销毁。 | |
| """ | |
| }, | |
| "label_content": { | |
| "en": "Content Image ", | |
| "zh": "内容图" | |
| }, | |
| "label_style": { | |
| "en": "Style Reference", | |
| "zh": "色彩参考图" | |
| }, | |
| "settings_header": { | |
| "en": "⚙️ Settings", | |
| "zh": "⚙️ 设置" | |
| }, | |
| "res_label": { | |
| "en": "Max Resolution Limit", | |
| "zh": "最大分辨率限制" | |
| }, | |
| "res_info": { | |
| "en": "Limits the input image size to speed up inference. 'Original' keeps original size.", | |
| "zh": "限制输入图像尺寸以加快推理速度。选择“Original”将保持原图尺寸。" | |
| }, | |
| "prot_label": { | |
| "en": "🧪 Enable Skin Protection (Experimental)", | |
| "zh": "🧪 肤色保护 (实验性功能)" | |
| }, | |
| "btn_run": { | |
| "en": "🚀 Start Color Transfer", | |
| "zh": "🚀 开始追色" | |
| }, | |
| "label_result": { | |
| "en": "Result Image", | |
| "zh": "追色结果" | |
| }, | |
| "example_label": { | |
| "en": "⚡ Quick Examples (Click to Load)", | |
| "zh": "⚡ 追色样例 (点击加载示例)" | |
| } | |
| } | |
| try: | |
| model_path = hf_hub_download( | |
| repo_id=model_repo_id, | |
| filename=model_filename, | |
| token=token_ | |
| ) | |
| except Exception as e: | |
| print(f"{e}") | |
| sess_options = ort.SessionOptions() | |
| sess_options.enable_cpu_mem_arena = False | |
| sess_options.intra_op_num_threads = 1 | |
| sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL | |
| providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] | |
| try: | |
| ort_session = ort.InferenceSession(model_path, providers=providers, sess_options=sess_options) | |
| input_names = [node.name for node in ort_session.get_inputs()] | |
| except Exception as e: | |
| print(f"{e}") | |
| ort_session = None | |
| try: | |
| script_path = hf_hub_download( | |
| repo_id=model_repo_id, | |
| filename=file_path, | |
| token=token_ | |
| ) | |
| skin_model_path = hf_hub_download( | |
| repo_id=model_repo_id, | |
| filename=pretrained_name, | |
| token=token_ | |
| ) | |
| spec = importlib.util.spec_from_file_location("skin_protection", script_path) | |
| secret_module = importlib.util.module_from_spec(spec) | |
| sys.modules["skin_protection"] = secret_module | |
| spec.loader.exec_module(secret_module) | |
| protector_instance = secret_module.SkinProtector(skin_model_path) | |
| print("✅ Skin protection module loaded successfully.") | |
| except Exception as e: | |
| print(f"⚠️ Skin protection not loaded (Running in standard mode): {e}") | |
| protector_instance = None | |
| def get_memory_mb(): | |
| pid = os.getpid() | |
| return psutil.Process(pid).memory_info().rss / 1024 / 1024 | |
| RESOLUTION_MAP = { | |
| "Original": 0, | |
| "4K": 3840 * 2160, # ~8.3MP | |
| "2K": 2560 * 1440, # ~3.7MP | |
| "1080p": 1920 * 1080, # ~2.1MP | |
| "720p": 1280 * 720 # ~0.9MP | |
| } | |
| def preprocess_image(img_rgb, resolution_choice="Original"): | |
| if img_rgb is None: | |
| return None | |
| max_pixel_limit = RESOLUTION_MAP.get(str(resolution_choice).split(" ")[0], 0) | |
| if max_pixel_limit > 0: | |
| h, w = img_rgb.shape[:2] | |
| num_pixels = h * w | |
| if num_pixels > max_pixel_limit: | |
| scale_factor = np.sqrt(max_pixel_limit / num_pixels) | |
| new_w = int(w * scale_factor) | |
| new_h = int(h * scale_factor) | |
| img_rgb = cv2.resize(img_rgb, (new_w, new_h), interpolation=cv2.INTER_AREA) | |
| # print(f"Resized image from {w}x{h} to {new_w}x{new_h}") | |
| img = img_rgb.astype(np.float32) / 255.0 | |
| img = img.transpose(2, 0, 1) | |
| img = np.expand_dims(img, axis=0) | |
| return img | |
| def postprocess_tensor(tensor): | |
| img = np.squeeze(tensor, axis=0) | |
| img = np.clip(img, 0, 1) | |
| img = img.transpose(1, 2, 0) | |
| img = (img * 255.0).astype(np.uint8) | |
| return img | |
| def run_inference_numpy(content_np, style_np): | |
| if ort_session is None: return content_np | |
| c_in = content_np.astype(np.float32) / 255.0 | |
| c_in = c_in.transpose(2, 0, 1)[np.newaxis, :, :, :] | |
| s_in = style_np.astype(np.float32) / 255.0 | |
| s_in = s_in.transpose(2, 0, 1)[np.newaxis, :, :, :] | |
| input_feed = {input_names[0]: c_in, input_names[1]: s_in} | |
| outputs = ort_session.run(None, input_feed) | |
| out_tensor = outputs[0] | |
| out_img = np.squeeze(out_tensor, axis=0) | |
| out_img = np.clip(out_img, 0, 1) | |
| out_img = out_img.transpose(1, 2, 0) | |
| out_img = (out_img * 255.0).astype(np.uint8) | |
| return out_img | |
| def inference(content_img, style_img, resolution_selection, enable_skin_protection): | |
| if ort_session is None: | |
| raise gr.Error("model do not exists") | |
| if content_img is None or style_img is None: | |
| return None | |
| # mem_start = get_memory_mb() | |
| # t_start = time.time() | |
| try: | |
| content_input = preprocess_image(content_img, resolution_selection) | |
| style_input = preprocess_image(style_img, resolution_selection) | |
| input_feed = { | |
| input_names[0]: content_input, | |
| input_names[1]: style_input | |
| } | |
| outputs = ort_session.run(None, input_feed) | |
| final_output_tensor = outputs[0] | |
| result_img = postprocess_tensor(final_output_tensor) | |
| if enable_skin_protection and protector_instance is not None: | |
| # print("skin protection activated") | |
| result_img = protector_instance.process( | |
| content_img=content_img, | |
| style_img=style_img, | |
| global_result_img=result_img, | |
| inference_callback=run_inference_numpy | |
| ) | |
| return result_img | |
| except Exception as e: | |
| print(f"{e}") | |
| raise gr.Error(f"{str(e)}") | |
| def toggle_language(current_lang): | |
| target_lang = "zh" if current_lang == "en" else "en" | |
| return ( | |
| target_lang, | |
| gr.update(value=I18N["lang_btn"][target_lang]), # Button | |
| gr.update(value=I18N["title"][target_lang]), # Markdown | |
| gr.update(label=I18N["desc_header"][target_lang]), # Accordion | |
| gr.update(value=I18N["desc_content"][target_lang]), # Markdown | |
| gr.update(label=I18N["label_content"][target_lang]), # Image Label | |
| gr.update(label=I18N["label_style"][target_lang]), # Image Label | |
| gr.update(value=I18N["btn_run"][target_lang]), # Button | |
| gr.update(label=I18N["label_result"][target_lang]), # Image Label | |
| gr.update(label=I18N["settings_header"][target_lang]), # Accordion Settings | |
| gr.update(label=I18N["res_label"][target_lang], info=I18N["res_info"][target_lang]), # Radio Label | |
| gr.update(label=I18N["prot_label"][target_lang]), | |
| ) | |
| custom_css = """ | |
| #col-container { | |
| margin: 0 auto; | |
| max-width: 1100px; | |
| } | |
| .gallery-container img{ | |
| object-fit: contain; | |
| } | |
| """ | |
| with gr.Blocks(css=custom_css, title="Color-Transfer") as demo: | |
| lang_state = gr.State("en") | |
| with gr.Column(elem_id="col-container"): | |
| with gr.Row(): | |
| with gr.Column(scale=5): | |
| md_title = gr.Markdown(I18N["title"]["en"]) | |
| with gr.Column(scale=0, min_width=80): | |
| btn_lang = gr.Button( | |
| value=I18N["lang_btn"]["en"], | |
| variant="secondary", | |
| size="sm" | |
| ) | |
| md_paper_code = gr.Markdown(I18N["paper_code"]["en"]) | |
| with gr.Accordion(label=I18N["desc_header"]["en"], open=False) as acc_desc: | |
| md_desc = gr.Markdown(I18N["desc_content"]["en"]) | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_content = gr.Image(label=I18N["label_content"]["en"], type="numpy", height=300, interactive=True) | |
| with gr.Column(): | |
| input_style = gr.Image(label=I18N["label_style"]["en"], type="numpy", height=300, interactive=True) | |
| with gr.Accordion(label=I18N["settings_header"]["en"], open=False) as acc_settings: | |
| res_choices = ["Original", "4K", "2K", "1080p", "720p"] | |
| radio_res = gr.Radio( | |
| choices=res_choices, | |
| value="Original", | |
| label=I18N["res_label"]["en"], | |
| info=I18N["res_info"]["en"], | |
| interactive=True | |
| ) | |
| with gr.Row(): | |
| chk_skin = gr.Checkbox( | |
| label=I18N["prot_label"]["en"], | |
| value=False, | |
| interactive=True | |
| ) | |
| with gr.Row(): | |
| btn_run = gr.Button(I18N["btn_run"]["en"], variant="primary") | |
| with gr.Row(): | |
| output_result = gr.Image(label=I18N["label_result"]["en"], type="numpy", interactive=False, height=450, format="png") | |
| btn_run.click( | |
| fn=lambda: None, | |
| inputs=None, | |
| outputs=output_result | |
| ).then( | |
| fn=inference, | |
| inputs=[input_content, input_style, radio_res, chk_skin], | |
| outputs=output_result | |
| ) | |
| gr.Examples( | |
| label=I18N["example_label"]["en"], | |
| examples=[ | |
| ["./figs/01_c.jpg", "./figs/01_s.jpg", "./figs/01.jpg"], | |
| ["./figs/02_c.jpg", "./figs/02_s.jpg", "./figs/02.jpg"], | |
| ["./figs/03_c.jpg", "./figs/03_s.jpg", "./figs/03.jpg"], | |
| ["./figs/04_c.jpg", "./figs/04_s.jpg", "./figs/04.png"], | |
| ], | |
| inputs=[input_content, input_style, output_result], | |
| ) | |
| btn_lang.click( | |
| fn=toggle_language, | |
| inputs=[lang_state], | |
| outputs=[ | |
| lang_state, | |
| btn_lang, | |
| md_title, | |
| acc_desc, | |
| md_desc, | |
| input_content, | |
| input_style, | |
| btn_run, | |
| output_result, | |
| acc_settings, | |
| radio_res, | |
| chk_skin, | |
| ] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |