Add examples with mask

app.py

@@ -107,16 +107,39 @@ async def process_image(
         return None
 
 
-def resize(image: Image.Image, shortest_side: int = 768) -> Image.Image:
-    if image.width <= shortest_side and image.height <= shortest_side:
-        return image
-    if image.width < image.height:
-        return image.resize(
-            size=(shortest_side, int(shortest_side * image.height / image.width))
-        )
-    return image.resize(
-        size=(int(shortest_side * image.width / image.height), shortest_side)
+def make_example(background_path: str, mask_path: str) -> EditorValue:
+    example1_background = np.array(Image.open(background_path))
+    example1_mask_only = np.array(Image.open(mask_path))[:, :, -1]
+
+    example1_layers = np.zeros(
+        (example1_background.shape[0], example1_background.shape[1], 4), dtype=np.uint8
+    )
+    example1_layers[:, :, 3] = example1_mask_only
+
+    example1_composite = np.zeros(
+        (example1_background.shape[0], example1_background.shape[1], 4), dtype=np.uint8
     )
+    example1_composite[:, :, :3] = example1_background
+    example1_composite[:, :, 3] = np.where(example1_mask_only == 255, 0, 255)
+
+    return {
+        "background": example1_background,
+        "layers": [example1_layers],
+        "composite": example1_composite,
+    }
+
+
+def resize_image(img: Image.Image, min_side_length: int = 768) -> Image.Image:
+    if img.width <= min_side_length and img.height <= min_side_length:
+        return img
+
+    aspect_ratio = img.width / img.height
+    if img.width < img.height:
+        new_height = int(min_side_length / aspect_ratio)
+        return img.resize((min_side_length, new_height))
+
+    new_width = int(min_side_length * aspect_ratio)
+    return img.resize((new_width, min_side_length))
 
 
 async def run_async(
@@ -124,24 +147,31 @@ async def run_async(
     progress: gr.Progress = gr.Progress(),
 ) -> tuple[Image.Image, Image.Image] | None:
     if not image_and_mask:
+        gr.Info("Please upload an image and draw a mask")
+        return None
+
+    image_np = image_and_mask["background"]
+    image_np = cast(np.ndarray, image_np)
+
+    # If the image is empty, return None
+    if np.sum(image_np) == 0:
+        gr.Info("Please upload an image")
         return None
 
     alpha_channel = image_and_mask["layers"][0]
     alpha_channel = cast(np.ndarray, alpha_channel)
     mask_np = np.where(alpha_channel[:, :, 3] == 0, 0, 255).astype(np.uint8)
 
-    image_np = image_and_mask["background"]
-    image_np = cast(np.ndarray, image_np)
+    # if mask_np is empty, return None
+    if np.sum(mask_np) == 0:
+        gr.Info("Please mark the areas you want to remove")
+        return None
 
-    # Save mask to ./masks.png
     mask = Image.fromarray(mask_np)
-    mask = resize(mask)
-    # mask.save("mask.png")
+    mask = resize_image(mask)
 
-    # Save image to ./images.png
     image = Image.fromarray(image_np)
-    image = resize(image)
-    # image.save("image.png")
+    image = resize_image(image)
 
     output = await process_image(
         image,  # type: ignore
@@ -150,9 +180,10 @@ async def run_async(
     )
 
     if output is None:
+        gr.Info("Processing failed")
         return None
 
-    return output, image
+    return image, output
 
 
 def run_sync(*args):
@@ -160,16 +191,31 @@ def run_sync(*args):
 
 
 with gr.Blocks() as demo:
-    gr.Markdown("""
-                # 🧹 Room Cleaner
-                Upload an image and and use pen tool (pencil icon at the bottom) to mark the areas you want to remove.
-                Click on the "Run" button to process the image and remove the marked areas.
-                """)
+    gr.HTML("""
+        <div style="text-align:center;">
+            <h1>🧹 Room Cleaner</h1>
+            <div>
+                <p>Upload an image and use the pencil tool (✏️ icon at the bottom) to <b>mark the areas you want to remove</b>.</p>
+                <p>For best results, include the shadows and reflections of the objects you want to remove.</p>
+                <p>You can remove multiple objects at once.</p>
+                <p>If you forget to mask some parts of your object, it's likely that the model will reconstruct them.</p>
+                <br>
+                <video width="640" height="360" controls style="margin: 0 auto; border-radius: 10px;">
+                    <source src="https://dropshare.blanchon.xyz/public/dropshare/room_cleaner_demo.mp4" type="video/mp4">
+                </video>
+                <br>
+                <p>Finally, click on the <b>"Run"</b> button to process the image.</p>
+                <p>Wait for the processing to complete and compare the original and processed images using the slider.</p>
+            
+                <p>⚠️ Note that the images are compressed to the workloads of the demo. </p>
+            </div>
+        </div>
+    """)
     with gr.Row():
         with gr.Column():
             # The image overflow, fix
             image_and_mask = gr.ImageMask(
-                label="Input Image and Mask",
+                label="Image and Mask",
                 layers=False,
                 show_fullscreen_button=False,
                 sources=["upload"],
@@ -177,11 +223,13 @@ with gr.Blocks() as demo:
                 interactive=True,
                 height="full",
                 width="full",
+                brush=gr.Brush(default_size=75, colors=["#000000"], color_mode="fixed"),
+                transforms=[],
             )
 
         with gr.Column():
             image_slider = ImageSlider(
-                label="Compare Original and Processed",
+                label="Result",
                 interactive=False,
             )
 
@@ -201,9 +249,13 @@ with gr.Blocks() as demo:
             )
 
     # Build examples
-    images_examples = glob.glob("examples/*")
-    mask_examples = [img.replace("inputs", "masks") for img in images_examples]
-    output_examples = [img.replace("inputs", "outputs") for img in images_examples]
+    images_examples = glob.glob("examples/*.jpg")
+    mask_examples = [
+        img.replace(".jpg", "") + "_mask_only.png" for img in images_examples
+    ]
+    output_examples = [
+        img.replace(".jpg", "") + "results.png" for img in images_examples
+    ]
     # examples = [
     #     [
     #         img,
@@ -212,50 +264,40 @@ with gr.Blocks() as demo:
     #     ]
     #     for img, mask, out in zip(images_examples, mask_examples, output_examples)
     # ]
+
+    example1 = make_example(images_examples[0], mask_examples[0])
+    example2 = make_example(images_examples[1], mask_examples[1])
+    example3 = make_example(images_examples[2], mask_examples[2])
+    example4 = make_example(images_examples[3], mask_examples[3])
+
     examples = [
         [
-            {
-                "background": "./examples/ex1.jpg",
-                "layers": [],
-                "composite": "./examples/ex1_mask.png",
-            },
-            # ("./examples/ex1.jpg", "./examples/ex1_result.png"),
+            example1,
+            # ("./examples/ex1.jpg", "./examples/ex1_result.png")
             (
                 "https://dropshare.blanchon.xyz/public/dropshare/ex1.jpg",
                 "https://dropshare.blanchon.xyz/public/dropshare/ex1_results.png",
             ),
         ],
         [
-            {
-                "background": "./examples/ex2.jpg",
-                "layers": [],
-                "composite": "./examples/ex2_mask.png",
-            },
-            # ("./examples/ex2.jpg", "./examples/ex2_result.png"),
+            example2,
+            # ("./examples/ex2.jpg", "./examples/ex2_result.png")
             (
                 "https://dropshare.blanchon.xyz/public/dropshare/ex2.jpg",
                 "https://dropshare.blanchon.xyz/public/dropshare/ex2_result.png",
             ),
         ],
         [
-            {
-                "background": "./examples/ex3.jpg",
-                "layers": [],
-                "composite": "./examples/ex3_mask.png",
-            },
-            # ("./examples/ex3.jpg", "./examples/ex3_result.png"),
+            example3,
+            # ("./examples/ex3.jpg", "./examples/ex3_result.png")
             (
                 "https://dropshare.blanchon.xyz/public/dropshare/ex3.jpg",
                 "https://dropshare.blanchon.xyz/public/dropshare/ex3_result.png",
             ),
         ],
         [
-            {
-                "background": "./examples/ex4.jpg",
-                "layers": [],
-                "composite": "./examples/ex4_mask.png",
-            },
-            # ("./examples/ex4.jpg", "./examples/ex4_result.png"),
+            example4,
+            # ("./examples/ex4.jpg", "./examples/ex4_result.png")
             (
                 "https://dropshare.blanchon.xyz/public/dropshare/ex4.jpg",
                 "https://dropshare.blanchon.xyz/public/dropshare/ex4_result.png",