initial test

app.py

@@ -0,0 +1,380 @@
+import gradio as gr
+import torch
+from PIL import Image, ImageFilter, ImageOps,ImageEnhance
+from scipy.ndimage import rank_filter, maximum_filter
+from skimage.filters import gabor
+import skimage.color
+import numpy as np
+import pathlib
+import glob
+import os
+from diffusers import StableDiffusionControlNetPipeline, DDIMScheduler, AutoencoderKL, ControlNetModel
+from ip_adapter import IPAdapter
+
+
+DESCRIPTION = """# [FilterPrompt](https://arxiv.org/abs/2404.13263): Guiding Imgae Transfer in Diffusion Models
+<img id="teaser" alt="teaser" src="https://raw.githubusercontent.com/Meaoxixi/FilterPrompt/gh-pages/resources/teaser.png" />
+"""
+# <img id="overview" alt="overview" src="https://github.com/Meaoxixi/FilterPrompt/blob/gh-pages/resources/teaser.png" />
+# 在你提供的链接中，你需要将 GitHub 页面链接中的 github.com 替换为 raw.githubusercontent.com，并将 blob 分支名称从链接中删除，以便获取原始文件内容。
+##################################################################################################################
+# 0. Get Pre-Models' Path Ready
+##################################################################################################################
+base_model_path = "models/stable-diffusion-v1-5"
+vae_model_path = "models/sd-vae-ft-mse"
+image_encoder_path = "models/IP-Adapter/image_encoder"
+ip_ckpt = "models/IP-Adapter/ip-adapter_sd15.bin"
+controlnet_softEdge_model_path = "models/ControlNet/ControlNet_depth"
+controlnet_depth_model_path = "models/ControlNet/ControlNet_softEdge"
+device = "cuda:0"
+##################################################################################################################
+# 1. load pipeline
+##################################################################################################################
+torch.cuda.empty_cache()
+## 1.1 noise_scheduler
+noise_scheduler = DDIMScheduler(
+        num_train_timesteps=1000,
+        beta_start=0.00085,
+        beta_end=0.012,
+        beta_schedule="scaled_linear",
+        clip_sample=False,
+        set_alpha_to_one=False,
+        steps_offset=1,
+    )
+#  1.2 vae
+vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
+#  1.3 ControlNet
+## 1.3.1 load controlnet_softEdge
+controlnet_softEdge = ControlNetModel.from_pretrained(controlnet_softEdge_model_path, torch_dtype=torch.float16)
+## 1.3.2 load controlnet_depth
+controlnet_depth = ControlNetModel.from_pretrained(controlnet_depth_model_path, torch_dtype=torch.float16)
+# 1.4 load SD pipeline
+pipe_softEdge = StableDiffusionControlNetPipeline.from_pretrained(
+        base_model_path,
+        controlnet=controlnet_softEdge,
+        torch_dtype=torch.float16,
+        scheduler=noise_scheduler,
+        vae=vae,
+        feature_extractor=None,
+        safety_checker=None
+    )
+pipe_depth = StableDiffusionControlNetPipeline.from_pretrained(
+        base_model_path,
+        controlnet=controlnet_depth,
+        torch_dtype=torch.float16,
+        scheduler=noise_scheduler,
+        vae=vae,
+        feature_extractor=None,
+        safety_checker=None
+    )
+print("1 Model loading completed !")
+print("##################################################################")
+def image_grid(imgs, rows, cols):
+    assert len(imgs) == rows * cols
+    w, h = imgs[0].size
+    grid = Image.new('RGB', size=(cols * w, rows * h))
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i % cols * w, i // cols * h))
+    return grid
+#########################################################################
+# 接下来是有关demo有关的函数定义
+## funcitions for task 1 : style transfer
+#########################################################################
+def gaussian_blur(image, blur_radius):
+    image = Image.open(image)
+    blurred_image = image.filter(ImageFilter.GaussianBlur(radius=blur_radius))
+    return blurred_image
+
+def task1_StyleTransfer(photo, blur_radius, sketch):
+    photoImage = Image.open(photo)
+    blurPhoto = gaussian_blur(photo, blur_radius)
+
+    Control_factor = 1.2
+    IP_factor = 0.6
+    ip_model = IPAdapter(pipe_depth, image_encoder_path, ip_ckpt, device, Control_factor=Control_factor, IP_factor=IP_factor)
+
+    depth_image= Image.open(sketch)
+    img_array = np.array(depth_image)
+    gray_img_array = np.dot(img_array[..., :3], [0.2989, 0.5870, 0.1140])
+    # 反相
+    inverted_array = 255 - gray_img_array
+    gray_img_array = inverted_array.astype(np.uint8)
+    processed_image = Image.fromarray(gray_img_array)
+    contrast_factor = 2
+    enhancer = ImageEnhance.Contrast(processed_image)
+    processed_image = enhancer.enhance(contrast_factor)
+
+    images = ip_model.generate(pil_image=photoImage, image=processed_image, num_samples=1, num_inference_steps=30, seed=52)
+    original = image_grid(images, 1, 1)
+    images = ip_model.generate(pil_image=blurPhoto, image=processed_image, num_samples=1, num_inference_steps=30, seed=52)
+    result= image_grid(images, 1, 1)
+
+    return original,result
+
+def task1_test(photo, blur_radius, sketch):
+    original = photo
+    print(type(original))
+    # <class 'str'>
+    result = sketch
+    return original, result
+#########################################################################
+## funcitions for task 2 : color transfer
+#########################################################################
+# 定义滤波器函数
+def desaturate_filter(image):
+    image = Image.open(image)
+    return ImageOps.grayscale(image)
+
+def gabor_filter(image):
+    image = Image.open(image)
+    image_array = np.array(image.convert('L'))  # 转换为灰度图像
+    filtered_real, filtered_imag = gabor(image_array, frequency=0.6)
+    filtered_image = np.sqrt(filtered_real**2 + filtered_imag**2)
+    return Image.fromarray(np.uint8(filtered_image))
+
+def rank_filter_func(image):
+    image = Image.open(image)
+    image_array = np.array(image.convert('L'))
+    filtered_image = rank_filter(image_array, rank=5, size=5)
+    return Image.fromarray(np.uint8(filtered_image))
+
+def max_filter_func(image):
+    image = Image.open(image)
+    image_array = np.array(image.convert('L'))
+    filtered_image = maximum_filter(image_array, size=20)
+    return Image.fromarray(np.uint8(filtered_image))
+# 定义处理函数
+def fun2(image,image2, filter_name):
+    if filter_name == "Desaturate Filter":
+        return desaturate_filter(image),desaturate_filter(image2)
+    elif filter_name == "Gabor Filter":
+        return gabor_filter(image),gabor_filter(image2)
+    elif filter_name == "Rank Filter":
+        return rank_filter_func(image),rank_filter_func(image2)
+    elif filter_name == "Max Filter":
+        return max_filter_func(image),max_filter_func(image2)
+    else:
+        return image,image2
+
+
+#############################################
+# Demo页面
+#############################################
+# 通过使用自定义 CSS 来控制风格
+#with gr.Blocks(theme=gr.themes.Soft()) as demo:
+with gr.Blocks(css="style.css") as demo:
+    # 0. 首先是标题和teaser
+    gr.Markdown(DESCRIPTION)
+
+    # 1. 第一个任务Style Transfer的界面代码（青铜器拓本转照片）
+    with gr.Group():
+        ## 1.1 任务描述
+        gr.Markdown(
+            """
+            ## Case 1: Style transfer
+                - In this task, our main goal is to achieve the style transfer from sketch to photo.
+                - In the original generation result, the surface of the object has redundant pattern representation from the style image.
+                - Next, you can control the Gaussian kernel size of GaussianBlur to weaken the expression of redundant pattern features in the generated results.
+            """)
+        ## 1.2 输入输出控件布局
+        #### 用Column()控制空间在列上的排列关系
+        with gr.Row():
+            # 第一列
+            with gr.Column():
+                with gr.Row():
+                    ### 1.2.1.1 输入真实照片
+                    photo = gr.Image(label="Input photo", type="filepath")
+                    print(photo)
+                    print(type(photo))
+                with gr.Row():
+                    ### 1.2.1.2 高斯核控件
+                    gaussianKernel = gr.Slider(minimum=0, maximum=8, step=1, value=2, label="Gaussian Blur Radius")
+            # 第二列
+            with gr.Column():
+                with gr.Row():
+                    # 1.2.2.1 输入素描图
+                    sketch = gr.Image(label="Input sketch", type="filepath")
+                    #print(sketch)
+                with gr.Row():
+                    # 1.2.2.2 按钮：开始生成图片
+                    task1Button = gr.Button("Preprocess")
+            # 第三列：显示初始的生成图
+            with gr.Column():
+                with gr.Row():
+                    original_result_task1 = gr.Image(label="Original generation result", interactive=False, type="pil")
+            # 第四列：显示使用高斯滤波之后的生成图
+            with gr.Column():
+                result_image_1 = gr.Image(label="Generate results after using GaussianBlur",type="pil")
+
+        ## 1.3 示例图展示
+        with gr.Row():
+            paths = sorted(pathlib.Path("images/inputExample").glob("*.jpg"))
+            gr.Examples(examples=[[path.as_posix()] for path in paths], inputs = sketch)
+        with gr.Row():
+            gr.Image(value="images/1_gaussian_filter.png", label=" Task example Image", type="filepath")
+
+    # 1. task 1 - style transfer 的界面代码写完了，现在写控件之间交互的逻辑
+    task1Button.click(
+        fn=task1_StyleTransfer,
+        #fn=task1_test,
+        inputs=[photo, gaussianKernel, sketch],
+        outputs=[original_result_task1, result_image_1],
+    )
+#
+# #  2. 第二个任务增强几何属性保护-Color transfer
+#     with gr.Group():
+#         ## 2.1 任务描述
+#         gr.Markdown(
+#             """
+#             ## Case 2: Color transferr
+#                 - In this task, our main goal is to transfer color from imageA to imageB. We can feel the effect of the filter on the protection of geometric properties.
+#                 - In the standard Controlnet-depth mode, the ideal input is the depth map.
+#                 - Here, we choose to input the result processed by some filters into the network instead of the original depth map.
+#                 - You can feel from the use of different filters that "decolorization+inversion+enhancement of contrast" can maximize the retention of detailed geometric information in the original image.
+#             """)
+#         ## 2.1 输入输出控件布局
+#         with gr.Row():
+#             with gr.Column():
+#                 with gr.Row():
+#                     input_appearImage = gr.Image(label="Input Appearance Image", type="filepath")
+#                 with gr.Row():
+#                     filter_dropdown = gr.Dropdown(
+#                         choices=["Desaturate Filter", "Gabor Filter", "Rank Filter", "Max Filter"],
+#                         label="Select Filter",
+#                         value="Desaturate Filter"
+#                     )
+#             with gr.Column():
+#                 with gr.Row():
+#                     input_strucImage = gr.Image(label="Input Structure Image", type="filepath")
+#                 with gr.Row():
+#                     geometry_button = gr.Button("Preprocess")
+#             with gr.Column():
+#                 with gr.Row():
+#                     afterFilterImage = gr.Image(label="Appearance image after filter choosed", interactive=False)
+#             with gr.Column():
+#                 result_task2 = gr.Image(label="Generate results")
+#                 # instyle = gr.State()
+#
+#         ## 2.3 示例图展示
+#         with gr.Row():
+#             gr.Image(value="task/color_transfer.png", label="example Image", type="filepath")
+#
+#         #  3. 第3个任务是光照效果的改善
+#         with gr.Group():
+#             ## 3.1 任务描述
+#             gr.Markdown(
+#                 """
+#                 ## Case 3: Image-to-Image translation
+#                     - In this example, our goal is to turn a simple outline drawing/sketch into a detailed and realistic photo.
+#                     - Here, we provide the original mask generation results, and provide the generation results after superimposing the image on the mask and passing the decolorization filter.
+#                     - From this, you can feel that the mask obtained by the decolorization operation can retain a certain amount of original lighting information and improve the texture of the generated results.
+#                     """)
+#             ## 3.2 输入输出控件布局
+#             with gr.Row():
+#                 with gr.Column():
+#                     with gr.Row():
+#                         input_appearImage = gr.Image(label="Input Appearance Image", type="filepath")
+#                     with gr.Row():
+#                         filter_dropdown = gr.Dropdown(
+#                             choices=["Desaturate Filter", "Gabor Filter", "Rank Filter", "Max Filter"],
+#                             label="Select Filter",
+#                             value="Desaturate Filter"
+#                         )
+#                 with gr.Column():
+#                     with gr.Row():
+#                         input_strucImage = gr.Image(label="Input Structure Image", type="filepath")
+#                     with gr.Row():
+#                         geometry_button = gr.Button("Preprocess")
+#                 with gr.Column():
+#                     with gr.Row():
+#                         afterFilterImage = gr.Image(label="Appearance image after filter choosed", interactive=False)
+#                 with gr.Column():
+#                     result_task2 = gr.Image(label="Generate results")
+#                     # instyle = gr.State()
+#
+#             ## 3.3 示例图展示
+#             with gr.Row():
+#                 gr.Image(value="task/4_light.jpg", label="example Image", type="filepath")
+#
+#         #  4. 第4个任务是materials transfer
+#         with gr.Group():
+#             ## 4.1 任务描述
+#             gr.Markdown(
+#                 """
+#                 ## Case 4: Materials Transfer
+#                     - In this example, our goal is to transfer the material appearance of one object image to another object image. The process involves changing the surface properties of objects in the image so that they appear to be made of another material.
+#                     - Here, we provide the original generation results and provide a variety of edited filters.
+#                     - You can specify any filtering operation and intuitively feel the impact of the filtering on the rendering properties in the generated results.
+#                     - For example, a sharpen filter can sharpen the texture of a stone, a Gaussian blur can smooth the texture of a stone, and a custom filter can change the style of a stone. These all show that filterPrompt is simple and intuitive.
+#                     """)
+#             ## 4.2 输入输出控件布局
+#             with gr.Row():
+#                 with gr.Column():
+#                     with gr.Row():
+#                         input_appearImage = gr.Image(label="Input Appearance Image", type="filepath")
+#                     with gr.Row():
+#                         filter_dropdown = gr.Dropdown(
+#                             choices=["Desaturate Filter", "Gabor Filter", "Rank Filter", "Max Filter"],
+#                             label="Select Filter",
+#                             value="Desaturate Filter"
+#                         )
+#                 with gr.Column():
+#                     with gr.Row():
+#                         input_strucImage = gr.Image(label="Input Structure Image", type="filepath")
+#                     with gr.Row():
+#                         geometry_button = gr.Button("Preprocess")
+#                 with gr.Column():
+#                     with gr.Row():
+#                         afterFilterImage = gr.Image(label="Appearance image after filter choosed", interactive=False)
+#                 with gr.Column():
+#                     result_task2 = gr.Image(label="Generate results")
+#                     # instyle = gr.State()
+#
+#             ## 3.3 示例图展示
+#             with gr.Row():
+#                 gr.Image(value="task/3mateialsTransfer.jpg", label="example Image", type="filepath")
+#
+#
+#     geometry_button.click(
+#         fn=fun2,
+#         inputs=[input_strucImage, input_appearImage, filter_dropdown],
+#         outputs=[afterFilterImage, result_task2],
+#     )
+    # aligned_face.change(
+    #     fn=model.reconstruct_face,
+    #     inputs=[aligned_face, encoder_type],
+    #     outputs=[
+    #         reconstructed_face,
+    #         instyle,
+    #     ],
+    # )
+    # style_type.change(
+    #     fn=update_slider,
+    #     inputs=style_type,
+    #     outputs=style_index,
+    # )
+    # style_type.change(
+    #     fn=update_style_image,
+    #     inputs=style_type,
+    #     outputs=style_image,
+    # )
+    # generate_button.click(
+    #     fn=model.generate,
+    #     inputs=[
+    #         style_type,
+    #         style_index,
+    #         structure_weight,
+    #         color_weight,
+    #         structure_only,
+    #         instyle,
+    #     ],
+    #     outputs=result,
+    #)
+##################################################################################################################
+# 2. run Demo on gradio
+##################################################################################################################
+
+if __name__ == "__main__":
+    demo.queue(max_size=5).launch()
+    #demo.queue(max_size=5).launch(server_port=12345)
+#demo.queue(max_size=5).launch(server_port=12345, share=True)