Update app.py

app.py

@@ -2,15 +2,21 @@ 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 = "selfie_multiclass_256x256.tflite"
+
+protector_instance = None
 
 I18N = {
     "lang_btn": {
@@ -68,8 +74,8 @@ I18N = {
         "zh": "色彩参考图"
     },
     "settings_header": {
-        "en": "⚙️ Settings (Resolution)",
-        "zh": "⚙️ 设置 (最大分辨率)"
+        "en": "⚙️ Settings",
+        "zh": "⚙️ 设置"
     },
     "res_label": {
         "en": "Max Resolution Limit",
@@ -79,6 +85,10 @@ I18N = {
         "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": "🚀 开始追色"
@@ -115,6 +125,30 @@ 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
@@ -157,7 +191,26 @@ def postprocess_tensor(tensor):
     img = (img * 255.0).astype(np.uint8)
     return img
 
-def inference(content_img, style_img, resolution_selection):
+
+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")
         
@@ -181,9 +234,14 @@ def inference(content_img, style_img, resolution_selection):
 
         result_img = postprocess_tensor(final_output_tensor)
 
-        # t_cost = time.time() - t_start
-        # mem_end = get_memory_mb()
-        # print(f"{t_cost:.3f}s | memory: {mem_start:.1f}MB -> {mem_end:.1f}MB")
+        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
 
@@ -206,6 +264,7 @@ def toggle_language(current_lang):
         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]),
     )
 
 
@@ -250,20 +309,26 @@ with gr.Blocks(css=custom_css, title="Color-Transfer") as demo:
                 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,          
+            fn=lambda: None,           
             inputs=None, 
             outputs=output_result
-        ).then(                   
+        ).then(                    
             fn=inference,
-            inputs=[input_content, input_style, radio_res],
+            inputs=[input_content, input_style, radio_res, chk_skin],
             outputs=output_result
         )
 
@@ -292,9 +357,10 @@ with gr.Blocks(css=custom_css, title="Color-Transfer") as demo:
                 input_content,  
                 input_style,  
                 btn_run,       
-                output_result,
-                acc_settings,
+                acc_settings, 
                 radio_res,
+                output_result,
+                chk_skin, 
             ]
         )