Update app.py (#1)

- Update app.py (be111dd97bcb3fba3def0a8384de64a5f118b7cd)


Co-authored-by: Aditya Deshmukh <[email protected]>

app.py

@@ -95,17 +95,14 @@ def get_masked_img(img, w, h, features, orig_h, orig_w, input_h, input_w, dilate
     return *figs, *masks
 
 
-def get_inpainted_img(img, mask0, mask1, mask2):
+def get_inpainted_img(img,mask):
     lama_config = args.lama_config
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    out = []
-    for mask in [mask0, mask1, mask2]:
-        if len(mask.shape)==3:
-            mask = mask[:,:,0]
-        img_inpainted = inpaint_img_with_builded_lama(
-            model['lama'], img, mask, lama_config, device=device)
-        out.append(img_inpainted)
-    return out
+    if len(mask.shape)==3:
+        mask = mask[:,:,0]
+    img_inpainted = inpaint_img_with_builded_lama(
+        model['lama'], img, mask, lama_config, device=device)
+    return img_inpainted
 
 
 # get args 
@@ -128,104 +125,16 @@ lama_ckpt = args.lama_ckpt
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model['lama'] = build_lama_model(lama_config, lama_ckpt, device=device)
 
-button_size = (100,50)
-with gr.Blocks() as demo:
-    features = gr.State(None)
-    orig_h = gr.State(None)
-    orig_w = gr.State(None)
-    input_h = gr.State(None)
-    input_w = gr.State(None)
-
-    with gr.Row().style(mobile_collapse=False, equal_height=True):
-        with gr.Column(variant="panel"):
-            with gr.Row():
-                gr.Markdown("## Input Image")
-            with gr.Row():
-                img = gr.Image(label="Input Image").style(height="200px")
-        with gr.Column(variant="panel"):
-            with gr.Row():
-                gr.Markdown("## Pointed Image")
-            with gr.Row():
-                img_pointed = gr.Plot(label='Pointed Image')
-        with gr.Column(variant="panel"):
-            with gr.Row():
-                gr.Markdown("## Control Panel")
-            with gr.Row():
-                w = gr.Number(label="Point Coordinate W")
-                h = gr.Number(label="Point Coordinate H")
-            dilate_kernel_size = gr.Slider(label="Dilate Kernel Size", minimum=0, maximum=100, step=1, value=15)
-            sam_mask = gr.Button("Predict Mask", variant="primary").style(full_width=True, size="sm")
-            lama = gr.Button("Inpaint Image", variant="primary").style(full_width=True, size="sm")
-            clear_button_image = gr.Button(value="Reset", label="Reset", variant="secondary").style(full_width=True, size="sm")
-
-    # todo: maybe we can delete this row, for it's unnecessary to show the original mask for customers
-    with gr.Row(variant="panel"):
-        with gr.Column():
-            with gr.Row():
-                gr.Markdown("## Segmentation Mask")
-            with gr.Row():
-                mask_0 = gr.outputs.Image(type="numpy", label="Segmentation Mask 0").style(height="200px")
-                mask_1 = gr.outputs.Image(type="numpy", label="Segmentation Mask 1").style(height="200px")
-                mask_2 = gr.outputs.Image(type="numpy", label="Segmentation Mask 2").style(height="200px")
-
-    with gr.Row(variant="panel"):
-        with gr.Column():
-            with gr.Row():
-                gr.Markdown("## Image with Mask")
-            with gr.Row():
-                img_with_mask_0 = gr.Plot(label="Image with Segmentation Mask 0")
-                img_with_mask_1 = gr.Plot(label="Image with Segmentation Mask 1")
-                img_with_mask_2 = gr.Plot(label="Image with Segmentation Mask 2")
-
-    with gr.Row(variant="panel"):
-        with gr.Column():
-            with gr.Row():
-                gr.Markdown("## Image Removed with Mask")
-            with gr.Row():
-                img_rm_with_mask_0 = gr.outputs.Image(
-                    type="numpy", label="Image Removed with Segmentation Mask 0").style(height="200px")
-                img_rm_with_mask_1 = gr.outputs.Image(
-                    type="numpy", label="Image Removed with Segmentation Mask 1").style(height="200px")
-                img_rm_with_mask_2 = gr.outputs.Image(
-                    type="numpy", label="Image Removed with Segmentation Mask 2").style(height="200px")
-
-
-    def get_select_coords(img, evt: gr.SelectData):
-        dpi = plt.rcParams['figure.dpi']
-        height, width = img.shape[:2]
-        fig = plt.figure(figsize=(width/dpi/0.77, height/dpi/0.77))
-        plt.imshow(img)
-        plt.axis('off')
-        plt.tight_layout()
-        show_points(plt.gca(), [[evt.index[0], evt.index[1]]], [1],
-                    size=(width*0.04)**2)
-        return evt.index[0], evt.index[1], fig
-
-    img.select(get_select_coords, [img], [w, h, img_pointed])
-    img.upload(get_sam_feat, [img], [features, orig_h, orig_w, input_h, input_w])
-
-    sam_mask.click(
-        get_masked_img,
-        [img, w, h, features, orig_h, orig_w, input_h, input_w, dilate_kernel_size],
-        [img_with_mask_0, img_with_mask_1, img_with_mask_2, mask_0, mask_1, mask_2]
-    )
-
-    lama.click(
-        get_inpainted_img,
-        [img, mask_0, mask_1, mask_2],
-        [img_rm_with_mask_0, img_rm_with_mask_1, img_rm_with_mask_2]
-    )
-
-
-    def reset(*args):
-        return [None for _ in args]
-
-    clear_button_image.click(
-        reset,
-        [img, features, img_pointed, w, h, mask_0, mask_1, mask_2, img_with_mask_0, img_with_mask_1, img_with_mask_2, img_rm_with_mask_0, img_rm_with_mask_1, img_rm_with_mask_2],
-        [img, features, img_pointed, w, h, mask_0, mask_1, mask_2, img_with_mask_0, img_with_mask_1, img_with_mask_2, img_rm_with_mask_0, img_rm_with_mask_1, img_rm_with_mask_2]
-    )
+image_input = gr.Image(label="Input Image")
+mask_input = gr.Image(label="Mask Image")
+demo = gr.Interface(
+    fn=get_inpainted_img,
+    inputs=[image_input, mask_input],
+    outputs=gr.Image(type="numpy", label="Output Image"),
+    title="Image and Mask Processor",
+    description="Upload an image and a mask to process the image. The mask highlights the areas to be processed.",
+)
 
 if __name__ == "__main__":
-    demo.launch()
-    
+    demo.queue(api_open=True, concurrency_count=2, max_size=10)
+    demo.launch(show_api=True)