Add flag save

app.py

@@ -1,7 +1,6 @@
 import csv
 import json
 import math
-import os
 import secrets
 from pathlib import Path
 from typing import cast
@@ -13,6 +12,7 @@ import torch
 from diffusers import FluxFillPipeline
 from gradio.components.gallery import GalleryMediaType
 from gradio.components.image_editor import EditorValue
+from huggingface_hub import HfApi
 from PIL import Image, ImageFilter, ImageOps
 
 DEVICE = "cuda"
@@ -25,38 +25,44 @@ SYSTEM_PROMPT = r"""This two-panel split-frame image showcases a furniture in as
 
 MASK_CONTEXT_PADDING = 16 * 8
 
-if os.environ.get("IN_SPACES", None):
-    FLAG_PATH = Path("/data/") / "flagged_data_points"
+api = HfApi()
+
+# Download the blanchon/FurnitureFlags init Path(__file__).parent / examples_dataset
+FLAG_PATH = Path(__file__).parent / "examples_dataset"
+if not torch.cuda.is_available():
+    FLAG_PATH = Path(__file__).parent / "examples_dataset"
+    FLAG_PATH.mkdir(parents=True, exist_ok=True)
 else:
-    FLAG_PATH = Path(__file__).parent / "flagged_data_points"
+    api.snapshot_download(
+        repo_id="blanchon/FurnitureFlags",
+        local_dir=FLAG_PATH,
+        repo_type="dataset",
+    )
 
 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,
-                ]
+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():
 
@@ -221,6 +227,13 @@ def flag(
     callback.flag(
         flag_data=[furniture_image_input, room_image_input, results],
     )
+    # Upload the flagged data points to the hub
+    api.upload_folder(
+        repo_id="blanchon/FurnitureFlags",
+        repo_type="dataset",
+        folder_path=FLAG_PATH,
+        ignore_patterns=[".cache"],
+    )
 
 
 @spaces.GPU(duration=150)
@@ -552,9 +565,9 @@ with gr.Blocks(css=css) as demo:
         [
             furniture_image_input,
             room_image_input,
-            # results,
+            results,
         ],
-        "flagged_data_points",
+        "examples_dataset",
     )
 
     with gr.Column(elem_id="col-showcase"):