update

app.py

@@ -230,8 +230,28 @@ class Tester(TesterBase):
             key=2
         )
 
+        st.markdown('<p class="big-font">Choose one mask for texture editing.</p>', unsafe_allow_html=True)
+        mask_list = glob(os.path.join("data/masks/*.png"))
+        byte_mask_list = []
+        for img_path in mask_list:
+            seg = Image.open(img_path).convert("L")
+            seg = np.asarray(seg)
+            seg = torch.from_numpy(seg).view(1, 1, seg.shape[0], seg.shape[1])
+            color_vq = self.draw_color_seg(seg)
+            vutils.save_image(color_vq, 'tmp/tmp.png')
+            with open('tmp/tmp.png', "rb") as image:
+                encoded = base64.b64encode(image.read()).decode()
+                byte_mask_list.append(f"data:image/jpeg;base64,{encoded}")
+        img_idx = clickable_images(
+            byte_mask_list,
+            titles=[f"Group #{str(i)}" for i in range(len(byte_mask_list))],
+            div_style={"display": "flex", "justify-content": "center", "flex-wrap": "wrap"},
+            img_style={"margin": "5px", "height": "150px"},
+        )
+        mask_path = mask_list[img_idx]
+
         st.markdown('<p class="big-font">Choose the texture segment for each group in the given mask below.</p>', unsafe_allow_html=True)
-        given_mask = Image.open('data/masks/124084_0_label.png').convert("L")
+        given_mask = Image.open(mask_path).convert("L")
         given_mask = np.asarray(given_mask)
         given_mask = torch.from_numpy(given_mask)
         H, W = given_mask.shape[0], given_mask.shape[1]