Add example images to the UI

.gitattributes

@@ -1 +1,5 @@
 *.pth filter=lfs diff=lfs merge=lfs -text
+*.webp filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -4,6 +4,7 @@ import numpy as np
 import torch
 import torch.nn.functional as F
 from skimage import img_as_ubyte
+from glob import glob
 
 from models import find_models_def
 
@@ -19,6 +20,17 @@ except Exception:  # pragma: no cover
 MODEL = None
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+EXAMPLES_DIR = os.path.join(BASE_DIR, "examples")
+
+
+def _list_example_images() -> list[str]:
+    if not os.path.isdir(EXAMPLES_DIR):
+        return []
+    exts = ("*.png", "*.jpg", "*.jpeg", "*.webp")
+    paths: list[str] = []
+    for ext in exts:
+        paths.extend(glob(os.path.join(EXAMPLES_DIR, ext)))
+    return sorted(paths)
 
 def _resolve_weights_path() -> str:
     """
@@ -138,6 +150,7 @@ demo = gr.Interface(
         "Space에 가중치(`checkpoints/model_best.pth`)가 포함되어 있거나, "
         "환경변수 WEIGHTS_REPO/WEIGHTS_FILENAME로 Hub에서 내려받을 수 있어야 합니다."
     ),
+    examples=_list_example_images(),
 )
 
 if __name__ == "__main__":