Add flag

app.py

@@ -1,4 +1,7 @@
+import csv
+import json
 import math
+import os
 import secrets
 from pathlib import Path
 from typing import cast
@@ -14,8 +17,6 @@ from PIL import Image, ImageFilter, ImageOps
 
 DEVICE = "cuda"
 
-EXAMPLES_DIR = Path(__file__).parent / "examples"
-
 MAX_SEED = np.iinfo(np.int32).max
 
 SYSTEM_PROMPT = r"""This two-panel split-frame image showcases a furniture in as a product shot versus styled in a room.
@@ -24,6 +25,39 @@ SYSTEM_PROMPT = r"""This two-panel split-frame image showcases a furniture in as
 
 MASK_CONTEXT_PADDING = 16 * 8
 
+if not os.environ.get("IN_SPACES", None):
+    FLAG_PATH = Path("/data/") / "flagged_data_points"
+else:
+    FLAG_PATH = Path(__file__).parent / "flagged_data_points"
+
+EXAMPLES: dict[str, list[str, str, str, list[str]]] = {}
+
+if not FLAG_PATH.exists():
+    FLAG_PATH.mkdir(parents=True)
+else:
+    flag_files = FLAG_PATH.glob("dataset*.csv")
+    for flag_file in flag_files:
+        with flag_file.open("r") as file:
+            reader = csv.reader(file)
+            next(reader)
+            for row in reader:
+                furniture_image, room_image, results_values, time = row
+                room_image = json.loads(room_image)
+                room_image_background = room_image["background"]
+                room_image_layers = room_image["layers"]
+                room_image_composite = room_image["composite"]
+                results_values = json.loads(results_values)
+                results_values = [result["image"] for result in results_values]
+                EXAMPLES[time] = [
+                    furniture_image,
+                    {
+                        "background": room_image_background,
+                        "layers": room_image_layers,
+                        "composite": room_image_composite,
+                    },
+                    # results_values,
+                ]
+
 if not torch.cuda.is_available():
 
     def _dummy_pipe(image: Image.Image, *args, **kwargs):  # noqa: ARG001
@@ -54,6 +88,9 @@ else:
     pipe.to(DEVICE)
 
 
+callback = gr.CSVLogger()
+
+
 def make_example(image_path: Path, mask_path: Path) -> EditorValue:
     background_image = Image.open(image_path)
     background_image = background_image.convert("RGB")
@@ -174,6 +211,18 @@ def adjust_bbox_to_divisible_16(
     return x_min, y_min, x_max, y_max
 
 
+def flag(
+    furniture_image_input: Image.Image,
+    room_image_input: EditorValue,
+    results: GalleryMediaType,
+):
+    if len(results) == 0:
+        return
+    callback.flag(
+        flag_data=[furniture_image_input, room_image_input, results],
+    )
+
+
 @spaces.GPU(duration=150)
 def infer(
     furniture_image_input: Image.Image,
@@ -376,18 +425,15 @@ with gr.Blocks(css=css) as demo:
                 max_height=50,
             )
             furniture_image_input = gr.Image(
-                label="Furniture Image",
+                label="furniture",
                 type="pil",
                 sources=["upload"],
                 image_mode="RGB",
                 height=500,
             )
             furniture_examples = gr.Examples(
-                examples=[
-                    EXAMPLES_DIR / "1" / "furniture_image.png",
-                    EXAMPLES_DIR / "2" / "furniture_image.png",
-                ],
-                examples_per_page=12,
+                examples=list({example[0] for example in EXAMPLES.values()}),
+                examples_per_page=6,
                 inputs=[furniture_image_input],
             )
         with gr.Column(elem_id="col-mid"):
@@ -402,7 +448,7 @@ with gr.Blocks(css=css) as demo:
                 max_height=50,
             )
             room_image_input = gr.ImageEditor(
-                label="Room Image - Draw mask for inpainting",
+                label="room_image",
                 type="pil",
                 sources=["upload"],
                 image_mode="RGBA",
@@ -411,16 +457,18 @@ with gr.Blocks(css=css) as demo:
                 height=500,
             )
             room_examples = gr.Examples(
-                examples=[
-                    make_example(
-                        EXAMPLES_DIR / "1" / "room_image.png",
-                        EXAMPLES_DIR / "1" / "room_mask.png",
-                    ),
-                    make_example(
-                        EXAMPLES_DIR / "2" / "room_image.png",
-                        EXAMPLES_DIR / "2" / "room_mask.png",
-                    ),
-                ],
+                examples=[example[1] for example in EXAMPLES.values()],
+                examples_per_page=6,
+                # examples=[
+                #     make_example(
+                #         EXAMPLES_DIR / "1" / "room_image.png",
+                #         EXAMPLES_DIR / "1" / "room_mask.png",
+                #     ),
+                #     make_example(
+                #         EXAMPLES_DIR / "2" / "room_image.png",
+                #         EXAMPLES_DIR / "2" / "room_mask.png",
+                #     ),
+                # ],
                 inputs=[room_image_input],
             )
         with gr.Column(elem_id="col-right"):
@@ -435,14 +483,15 @@ with gr.Blocks(css=css) as demo:
                 max_height=50,
             )
             results = gr.Gallery(
-                label="Results",
+                label="results",
                 show_label=False,
                 columns=2,
                 height=500,
-                interactive=False,
+                format="png",
+                # interactive=False,
             )
             run_button = gr.Button("Run")
-
+            flag_button = gr.Button("Flag")
             # Reset the results when the run button is clicked
             run_button.click(
                 outputs=results,
@@ -498,6 +547,16 @@ with gr.Blocks(css=css) as demo:
                         value=20,
                     )
 
+    # This needs to be called at some point prior to the first call to callback.flag()
+    callback.setup(
+        [
+            furniture_image_input,
+            room_image_input,
+            # results,
+        ],
+        "flagged_data_points",
+    )
+
     with gr.Column(elem_id="col-showcase"):
         gr.HTML("""
         <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
@@ -509,28 +568,13 @@ with gr.Blocks(css=css) as demo:
         </div>
         """)
         show_case = gr.Examples(
-            examples=[
-                [
-                    EXAMPLES_DIR / "1" / "furniture_image.png",
-                    make_example(
-                        EXAMPLES_DIR / "1" / "room_image.png",
-                        EXAMPLES_DIR / "1" / "room_mask.png",
-                    ),
-                ],
-                [
-                    EXAMPLES_DIR / "2" / "furniture_image.png",
-                    make_example(
-                        EXAMPLES_DIR / "2" / "room_image.png",
-                        EXAMPLES_DIR / "2" / "room_mask.png",
-                    ),
-                ],
-            ],
-            inputs=[furniture_image_input, room_image_input],
+            examples=list(EXAMPLES.values()),
+            inputs=[furniture_image_input, room_image_input, results],
             label=None,
+            examples_per_page=12,
         )
 
-    gr.on(
-        triggers=[run_button.click],
+    run_button.click(
         fn=infer,
         inputs=[
             furniture_image_input,
@@ -545,5 +589,10 @@ with gr.Blocks(css=css) as demo:
         ],
         outputs=[results, seed],
     )
+    flag_button.click(
+        fn=flag,
+        inputs=[furniture_image_input, room_image_input, results],
+        preprocess=False,
+    )
 
 demo.launch()