Update

app.py

@@ -36,12 +36,16 @@ imgs_tensor = torch.zeros(4, 3, patch_h * 14, patch_w * 14)
 # PCA
 pca = PCA(n_components=3)
 
-def query_image(img1, img2, img3, img4) -> List[np.ndarray]:
+def query_image(
+    img1, img2, img3, img4,
+    background_threshold,
+    is_foreground_larger_than_threshold,
+) -> List[np.ndarray]:
 
     # Transform
     imgs = [img1, img2, img3, img4]
     for i, img in enumerate(imgs):
-        img = np.transpose(img, (2, 0, 1))
+        img = np.transpose(img, (2, 0, 1)) / 255
         imgs_tensor[i] = transform(torch.Tensor(img))
 
     # Get feature from patches
@@ -56,7 +60,10 @@ def query_image(img1, img2, img3, img4) -> List[np.ndarray]:
     pca_feature = sklearn.preprocessing.minmax_scale(pca_features)
 
     # Foreground/Background
-    pca_features_bg = pca_features[:, 0] < 0
+    if is_foreground_larger_than_threshold:
+        pca_features_bg = pca_features[:, 0] < background_threshold
+    else:
+        pca_features_bg = pca_features[:, 0] > background_threshold
     pca_features_fg = ~pca_features_bg
 
     # PCA with only foreground
@@ -79,10 +86,14 @@ DINOV2 PCA
 """
 demo = gr.Interface(
     query_image,
-    inputs=[gr.Image(), gr.Image(), gr.Image(), gr.Image()],
+    inputs=[gr.Image(), gr.Image(), gr.Image(), gr.Image(), gr.Slider(-1, 1, value=0.1), gr.Checkbox(label="foreground is larger than threshold", value=True) ],
     outputs=[gr.Image(), gr.Image(), gr.Image(), gr.Image()],
     title="DINOV2 PCA",
     description=description,
-    examples=[],
+    examples=[
+        ["assets/1.png", "assets/2.png","assets/3.png","assets/4.png", 0.9, True],
+        ["assets/5.png", "assets/6.png","assets/7.png","assets/8.png", 0.6, True],
+        ["assets/9.png", "assets/10.png","assets/11.png","assets/12.png", 0.6, True],
+    ]
 )
 demo.launch()