Update app.py

app.py

@@ -66,27 +66,54 @@ def resize_image(input_image, resolution):
     img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
     return img
 
-def process(original_image, input_mask, prompt, negative_prompt, blended, invert_mask, control_strength, seed, randomize_seed, guidance_scale, num_inference_steps):
+
+def process(input_image, 
+    original_image, 
+    original_mask, 
+    input_mask, 
+    selected_points, 
+    prompt, 
+    negative_prompt, 
+    blended, 
+    invert_mask, 
+    control_strength, 
+    seed, 
+    randomize_seed, 
+    guidance_scale, 
+    num_inference_steps):
     if original_image is None:
         raise gr.Error('Please upload the input image')
-    if input_mask is None:
-        raise gr.Error("Please upload a white-black Mask image")
+    if (original_mask is None or len(selected_points)==0) and input_mask is None:
+        raise gr.Error("Please click the region where you hope unchanged/changed, or upload a white-black Mask image")
     
-    H, W = original_image.shape[:2]
-    original_mask = cv2.resize(input_mask, (W, H))
+    # load example image
+    if isinstance(original_image, int):
+        image_name = image_examples[original_image][0]
+        original_image = cv2.imread(image_name)
+        original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
+
+    if input_mask is not None:
+        H,W=original_image.shape[:2]
+        original_mask = cv2.resize(input_mask, (W, H))
+    else:
+        original_mask = np.clip(255 - original_mask, 0, 255).astype(np.uint8)
 
     if invert_mask:
-        original_mask = 255 - original_mask
-    mask = 1.*(original_mask.sum(-1) > 255)[:,:,np.newaxis]
-    masked_image = original_image * (1 - mask)
+        original_mask=255-original_mask
+
+    mask = 1.*(original_mask.sum(-1)>255)[:,:,np.newaxis]
+    masked_image = original_image * (1-mask)
+
     init_image = Image.fromarray(masked_image.astype(np.uint8)).convert("RGB")
     mask_image = Image.fromarray(original_mask.astype(np.uint8)).convert("RGB")
-    generator = torch.Generator("cuda").manual_seed(random.randint(0, 2147483647) if randomize_seed else seed)
+
+    generator = torch.Generator("cuda").manual_seed(random.randint(0,2147483647) if randomize_seed else seed)
+
     image = pipe(
-        [prompt]*2,
-        init_image,
-        mask_image,
-        num_inference_steps=num_inference_steps,
+        [prompt]*2, 
+        init_image, 
+        mask_image, 
+        num_inference_steps=num_inference_steps, 
         guidance_scale=guidance_scale,
         generator=generator,
         brushnet_conditioning_scale=float(control_strength),
@@ -94,43 +121,220 @@ def process(original_image, input_mask, prompt, negative_prompt, blended, invert
     ).images
 
     if blended:
-        if control_strength < 1.0:
+        if control_strength<1.0:
             raise gr.Error('Using blurred blending with control strength less than 1.0 is not allowed')
-        blended_image = []
+        blended_image=[]
+        # blur, you can adjust the parameters for better performance
         mask_blurred = cv2.GaussianBlur(mask*255, (21, 21), 0)/255
         mask_blurred = mask_blurred[:,:,np.newaxis]
-        mask = 1 - (1 - mask) * (1 - mask_blurred)
+        mask = 1-(1-mask) * (1-mask_blurred)
         for image_i in image:
-            image_np = np.array(image_i)
-            image_pasted = original_image * (1 - mask) + image_np * mask
-            image_pasted = image_pasted.astype(image_np.dtype)
+            image_np=np.array(image_i)
+            image_pasted=original_image * (1-mask) + image_np*mask
+
+            image_pasted=image_pasted.astype(image_np.dtype)
             blended_image.append(Image.fromarray(image_pasted))
-        image = blended_image
+        
+        image=blended_image
 
     return image
 
-# Create Gradio interface
-with gr.Blocks() as demo:
+block = gr.Blocks(
+        theme=gr.themes.Soft(
+             radius_size=gr.themes.sizes.radius_none,
+             text_size=gr.themes.sizes.text_md
+         )
+        ).queue()
+with block:
     with gr.Row():
         with gr.Column():
-            original_image = gr.Image(type="numpy", label="Original Image")
-            input_mask = gr.Image(type="numpy", label="Mask Image")
-            prompt = gr.Textbox(label="Prompt")
-            negative_prompt = gr.Textbox(label="Negative Prompt", value='ugly, low quality')
-            blended = gr.Checkbox(label="Blurred Blending", value=False)
-            invert_mask = gr.Checkbox(label="Invert Mask", value=False)
-            control_strength = gr.Slider(label="Control Strength", minimum=0, maximum=1.1, value=1, step=0.01)
-            seed = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, value=551793204)
-            randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
-            guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=12, step=0.1, value=7.5)
-            num_inference_steps = gr.Slider(label="Number of Inference Steps", minimum=1, maximum=50, step=1, value=50)
-            run_button = gr.Button("Run")
-        
+            
+            gr.HTML(f"""
+                    <div style="text-align: center;">
+                        <h1>BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion</h1>
+                        <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
+                            <a href=""></a>
+                            <a href='https://tencentarc.github.io/BrushNet/'><img src='https://img.shields.io/badge/Project_Page-BrushNet-green' alt='Project Page'></a>
+                            <a href='https://arxiv.org/abs/2403.06976'><img src='https://img.shields.io/badge/Paper-Arxiv-blue'></a>
+                        </div>
+                        </br>
+                    </div>
+            """)
+
+
+    with gr.Accordion(label="🧭 Instructions:", open=True, elem_id="accordion"):
+        with gr.Row(equal_height=True):
+            gr.Markdown("""
+            - ⭐️ <b>step1: </b>Upload or select one image from Example
+            - ⭐️ <b>step2: </b>Click on Input-image to select the object to be retained (or upload a white-black Mask image, in which white color indicates the region you want to keep unchanged). You can tick the 'Invert Mask' box to switch region unchanged and change.
+            - ⭐️ <b>step3: </b>Input prompt for generating new contents
+            - ⭐️ <b>step4: </b>Click Run button
+            """)                          
+    with gr.Row():
+        with gr.Column():
+            with gr.Column(elem_id="Input"):
+                with gr.Row():
+                    with gr.Tabs(elem_classes=["feedback"]):
+                        with gr.TabItem("Input Image"):
+                            input_image = gr.Image(type="numpy", label="input",scale=2, height=640)
+                original_image = gr.State(value=None,label="index")
+                original_mask = gr.State(value=None)
+                selected_points = gr.State([],label="select points")
+                with gr.Row(elem_id="Seg"):
+                    radio = gr.Radio(['foreground', 'background'], label='Click to seg: ', value='foreground',scale=2)
+                    undo_button = gr.Button('Undo seg', elem_id="btnSEG",scale=1)
+            prompt = gr.Textbox(label="Prompt", placeholder="Please input your prompt",value='',lines=1)
+            negative_prompt = gr.Text(
+                        label="Negative Prompt",
+                        max_lines=5,
+                        placeholder="Please input your negative prompt",
+                        value='ugly, low quality',lines=1
+                    )
+            with gr.Group():
+                with gr.Row():
+                    blending = gr.Checkbox(label="Blurred Blending", value=False)
+                    invert_mask = gr.Checkbox(label="Invert Mask", value=True)
+            run_button = gr.Button("Run",elem_id="btn")
+            
+            with gr.Accordion("More input params (highly-recommended)", open=False, elem_id="accordion1"):
+                control_strength = gr.Slider(
+                    label="Control Strength: ", show_label=True, minimum=0, maximum=1.1, value=1, step=0.01
+                    )
+                with gr.Group():
+                    seed = gr.Slider(
+                        label="Seed: ", minimum=0, maximum=2147483647, step=1, value=551793204
+                    )
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
+                
+                with gr.Group():
+                    with gr.Row():
+                        guidance_scale = gr.Slider(
+                            label="Guidance scale",
+                            minimum=1,
+                            maximum=12,
+                            step=0.1,
+                            value=12,
+                        )
+                        num_inference_steps = gr.Slider(
+                            label="Number of inference steps",
+                            minimum=1,
+                            maximum=50,
+                            step=1,
+                            value=50,
+                        )
+                with gr.Row(elem_id="Image"):
+                    with gr.Tabs(elem_classes=["feedback1"]):
+                        with gr.TabItem("User-specified Mask Image (Optional)"):
+                            input_mask = gr.Image(type="numpy", label="Mask Image", height=640)
+            
         with gr.Column():
-            result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", preview=True)
+            with gr.Tabs(elem_classes=["feedback"]):
+                with gr.TabItem("Outputs"):
+                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", preview=True)
+    with gr.Row():
+        def process_example(input_image, prompt, input_mask, original_image, selected_points,result_gallery): #
+            return input_image, prompt, input_mask, original_image, [], result_gallery
+        example = gr.Examples(
+            label="Input Example",
+            examples=image_examples,
+            inputs=[input_image, prompt, input_mask, original_image, selected_points,result_gallery],
+            outputs=[input_image, prompt, input_mask, original_image, selected_points],
+            fn=process_example,
+            run_on_click=True,
+            examples_per_page=10
+        )
+
+    # once user upload an image, the original image is stored in `original_image`
+    def store_img(img):
+        # image upload is too slow
+        if min(img.shape[0], img.shape[1]) > 512:
+            img = resize_image(img, 512)
+        if max(img.shape[0], img.shape[1])*1.0/min(img.shape[0], img.shape[1])>2.0:
+            raise gr.Error('image aspect ratio cannot be larger than 2.0')
+        return img, img, [], None  # when new image is uploaded, `selected_points` should be empty
+
+    input_image.upload(
+        store_img,
+        [input_image],
+        [input_image, original_image, selected_points]
+    )
+
+    # user click the image to get points, and show the points on the image
+    def segmentation(img, sel_pix):
+        # online show seg mask
+        points = []
+        labels = []
+        for p, l in sel_pix:
+            points.append(p)
+            labels.append(l)
+        mobile_predictor.set_image(img if isinstance(img, np.ndarray) else np.array(img))
+        with torch.no_grad():
+            masks, _, _ = mobile_predictor.predict(point_coords=np.array(points), point_labels=np.array(labels), multimask_output=False)
+
+        output_mask = np.ones((masks.shape[1], masks.shape[2], 3))*255
+        for i in range(3):
+                output_mask[masks[0] == True, i] = 0.0
+
+        mask_all = np.ones((masks.shape[1], masks.shape[2], 3))
+        color_mask = np.random.random((1, 3)).tolist()[0]
+        for i in range(3):
+            mask_all[masks[0] == True, i] = color_mask[i]
+        masked_img = img / 255 * 0.3 + mask_all * 0.7
+        masked_img = masked_img*255
+        ## draw points
+        for point, label in sel_pix:
+            cv2.drawMarker(masked_img, point, colors[label], markerType=markers[label], markerSize=20, thickness=5)
+        return masked_img, output_mask
+    
+    def get_point(img, sel_pix, point_type, evt: gr.SelectData):
+        if point_type == 'foreground':
+            sel_pix.append((evt.index, 1))   # append the foreground_point
+        elif point_type == 'background':
+            sel_pix.append((evt.index, 0))    # append the background_point
+        else:
+            sel_pix.append((evt.index, 1))    # default foreground_point
+
+        if isinstance(img, int):
+            image_name = image_examples[img][0]
+            img = cv2.imread(image_name)
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+        # online show seg mask
+        masked_img, output_mask = segmentation(img, sel_pix)
+        return masked_img.astype(np.uint8), output_mask
+    
+    input_image.select(
+        get_point,
+        [original_image, selected_points, radio],
+        [input_image, original_mask],
+    )
+
+    # undo the selected point
+    def undo_points(orig_img, sel_pix):
+        # draw points
+        output_mask = None
+        if len(sel_pix) != 0:
+            if isinstance(orig_img, int):   # if orig_img is int, the image if select from examples
+                temp = cv2.imread(image_examples[orig_img][0])
+                temp = cv2.cvtColor(temp, cv2.COLOR_BGR2RGB)
+            else:
+                temp = orig_img.copy()
+            sel_pix.pop()
+            # online show seg mask
+            if len(sel_pix) !=0:
+                temp, output_mask = segmentation(temp, sel_pix)
+            return temp.astype(np.uint8), output_mask
+        else:
+            gr.Error("Nothing to Undo")
+    
+    undo_button.click(
+        undo_points,
+        [original_image, selected_points],
+        [input_image, original_mask]
+    )
+
+    ips=[input_image, original_image, original_mask, input_mask, selected_points, prompt, negative_prompt, blending, invert_mask, control_strength, seed, randomize_seed, guidance_scale, num_inference_steps]
+    run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
 
-    inputs = [original_image, input_mask, prompt, negative_prompt, blended, invert_mask, control_strength, seed, randomize_seed, guidance_scale, num_inference_steps]
-    run_button.click(fn=process, inputs=inputs, outputs=[result_gallery])
 
-demo.queue(concurrency_count=1, max_size=1, api_open=True)
-demo.launch(show_api=True)
+block.launch()