add filter

app_base.py

@@ -5,7 +5,7 @@ import torch
 import os
 import gc
 
-from PIL import Image, ImageEnhance
+from PIL import Image, ImageEnhance, ImageFilter
 from segment_utils import(
     segment_image,
     restore_result_and_save,
@@ -18,6 +18,16 @@ DEFAULT_EDIT_PROMPT = "a person with perfect face"
 
 DEFAULT_CATEGORY = "face"
 
+filter_names = [
+    "NONE",
+    "DETAIL",
+    "SMOOTH",
+    "SMOOTH_MORE",
+    "SHARPEN",
+    "EDGE_ENHANCE",
+    "EDGE_ENHANCE_MORE",
+]
+
 @spaces.GPU(duration=10)
 @torch.inference_mode()
 @torch.no_grad()
@@ -31,6 +41,10 @@ def image_to_image(
     start_step: int,
     guidance_scale: float,
     brightness: float = 1.0,
+    color: float = 1.0,
+    contrast: float = 1.0,
+    sharpness: float = 1.0,
+    filter: str = "NONE",
 ):
     w2 = 1.0
     run_task_time = 0
@@ -70,6 +84,29 @@ def image_to_image(
 
     enhancer = ImageEnhance.Brightness(enhanced_image)
     enhanced_image = enhancer.enhance(brightness)
+    enhancer = ImageEnhance.Color(enhanced_image)
+    enhanced_image = enhancer.enhance(color)
+    enhancer = ImageEnhance.Contrast(enhanced_image)
+    enhanced_image = enhancer.enhance(contrast)
+    enhancer = ImageEnhance.Sharpness(enhanced_image)
+    enhanced_image = enhancer.enhance(sharpness)
+
+    if filter == "NONE":
+        pass
+    elif filter == "DETAIL":
+        enhanced_image = enhanced_image.filter(ImageFilter.DETAIL)
+    elif filter == "SMOOTH":
+        enhanced_image = enhanced_image.filter(ImageFilter.SMOOTH)
+    elif filter == "SMOOTH_MORE":
+        enhanced_image = enhanced_image.filter(ImageFilter.SMOOTH_MORE)
+    elif filter == "SHARPEN":
+        enhanced_image = enhanced_image.filter(ImageFilter.SHARPEN)
+    elif filter == "EDGE_ENHANCE":
+        enhanced_image = enhanced_image.filter(ImageFilter.EDGE_ENHANCE)
+    elif filter == "EDGE_ENHANCE_MORE":
+        enhanced_image = enhanced_image.filter(ImageFilter.EDGE_ENHANCE_MORE)
+
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'image_enhance done')
 
     return enhanced_image, time_cost_str
 
@@ -108,7 +145,11 @@ def create_demo() -> gr.Blocks:
             with gr.Column():
                 input_image_prompt = gr.Textbox(lines=1, label="Input Image Prompt", value=DEFAULT_SRC_PROMPT)
                 edit_prompt = gr.Textbox(lines=1, label="Edit Prompt", value=DEFAULT_EDIT_PROMPT)
-                brightness = gr.Slider(minimum=0, maximum=2, value=1.0, step=0.05, label="Brightness")
+                with gr.Row():
+                    brightness = gr.Number(label="Brightness", value=1.0, minimum=0.0, maximum=2.0, step=0.01)
+                    color = gr.Number(label="Color", value=1.0, minimum=0.0, maximum=2.0, step=0.01)
+                    contrast = gr.Number(label="Contrast", value=1.0, minimum=0.0, maximum=2.0, step=0.01)
+                    sharpness = gr.Number(label="Sharpness", value=1.0, minimum=0.0, maximum=2.0, step=0.01)
                 with gr.Accordion("Advanced Options", open=False):
                     category = gr.Textbox(label="Category", value=DEFAULT_CATEGORY, visible=False)
                     mask_expansion = gr.Number(label="Mask Expansion", value=50, visible=True)
@@ -117,6 +158,7 @@ def create_demo() -> gr.Blocks:
             with gr.Column():
                 num_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Num Steps")
                 start_step = gr.Slider(minimum=1, maximum=100, value=15, step=1, label="Start Step")
+                filter = gr.Dropdown(choices=filter_names, label="Filter", value="NONE")
                 g_btn = gr.Button("Edit Image")
                 with gr.Accordion("Advanced Options", open=False):
                     guidance_scale = gr.Slider(minimum=0, maximum=20, value=0, step=0.5, label="Guidance Scale")
@@ -140,7 +182,7 @@ def create_demo() -> gr.Blocks:
             outputs=[origin_area_image, cropper],
         ).success(
             fn=image_to_image,
-            inputs=[origin_area_image, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step, guidance_scale, brightness],
+            inputs=[origin_area_image, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step, guidance_scale, brightness, color, contrast, sharpness, filter],
             outputs=[enhanced_image, generated_cost],
         ).success(
             fn=restore_result_and_save,