Update app.py

app.py

@@ -143,6 +143,25 @@ def get_mask(image, label):
 
     return result_mask
 
+def create_highlighted_mask(image_np, mask_np, alpha=0.5, gray_value=128):
+
+
+    if mask_np.max() <= 1.0:
+        mask_np = (mask_np * 255).astype(np.uint8)
+    mask_bool = mask_np > 128
+
+    image_float = image_np.astype(np.float32)
+
+    # 灰色图层
+    gray_overlay = np.full_like(image_float, gray_value, dtype=np.float32)
+
+    # 混合
+    result = image_float.copy()
+    result[mask_bool] = (
+        (1 - alpha) * image_float[mask_bool] + alpha * gray_overlay[mask_bool]
+    )
+
+    return result.astype(np.uint8)
 
 hf_token = os.getenv("HF_TOKEN")
 
@@ -278,6 +297,8 @@ def run_local(base_image, base_mask, reference_image, ref_mask, seed, base_mask_
     tar_image = cv2.resize(tar_image, size)
     diptych_ref_tar = np.concatenate([masked_ref_image, tar_image], axis=1)
 
+    show_diptych_ref_tar = create_highlighted_mask(diptych_ref_tar, mask_diptych)
+    show_diptych_ref_tar = Image.fromarray(show_diptych_ref_tar)
 
     tar_mask = np.stack([tar_mask,tar_mask,tar_mask],-1)
     mask_black = np.ones_like(tar_image) * 0
@@ -316,9 +337,9 @@ def run_local(base_image, base_mask, reference_image, ref_mask, seed, base_mask_
     edited_image = Image.fromarray(edited_image)
 
     if ref_mask_option != "Label to Mask":
-        return [edited_image]
+        return [edited_image, show_diptych_ref_tar]
     else:
-        return [return_ref_mask, edited_image]   
+        return [return_ref_mask, show_diptych_ref_tar, edited_image]   
 
 def update_ui(option):
     if option == "Draw Mask":
@@ -357,15 +378,18 @@ with gr.Blocks() as demo:
                 ref_mask_option = gr.Radio(["Draw Mask", "Upload with Mask", "Label to Mask"], label="Reference Mask Input Option", value="Upload with Mask")
 
             with gr.Row():
-                text_prompt = gr.Textbox(label="Label")
+                text_prompt = gr.Textbox(label="Label", placeholder="Enter the category of the reference object, e.g., car, dress, toy, etc.")
 
         with gr.Column(scale=1):
-            baseline_gallery = gr.Gallery(label='Output', show_label=True, elem_id="gallery", height=765, columns=1)
+            baseline_gallery = gr.Gallery(label='Output', show_label=True, elem_id="gallery", height=705, columns=1)
             with gr.Accordion("Advanced Option", open=True):
                 seed = gr.Slider(label="Seed", minimum=-1, maximum=999999999, step=1, value=666)
                 gr.Markdown("### Guidelines")
                 gr.Markdown(" Users can try using different seeds. For example, seeds like 42 and 123456 may produce different effects.")
-                gr.Markdown(" Label to Mask means generating a mask by simply inputting a label.")
+                gr.Markdown(" Draw Mask means manually drawing a mask on the original image.")
+                gr.Markdown(" Upload with Mask means uploading a mask file.")
+                gr.Markdown(" Label to Mask means simply inputting a label to automatically extract the mask and obtain the result.")
+                
 
     run_local_button = gr.Button(value="Run")