Update app.py

app.py

@@ -69,10 +69,10 @@ class Tester(TesterBase):
     def to_pil(self, tensor):
         return transforms.ToPILImage()(tensor.cpu().squeeze().clamp(0.0, 1.0)).convert("RGB")
 
-    def display(self):
+    def display_synthesis(self):
         with st.spinner('Running...'):
             with torch.no_grad():
-                grouping_mask = self.model_forward(self.data, self.slic, return_type = 'grouping')
+                grouping_mask = self.model_forward_synthesis(self.data, self.slic, return_type = 'grouping')
 
             data = (self.data + 1) / 2.0
 
@@ -124,7 +124,7 @@ class Tester(TesterBase):
                 tex_size = st.slider('', 0, 1000, 256)
                 tex_size = (tex_size // 8) * 8
                 with torch.no_grad():
-                    tex = self.model_forward(self.data, self.slic, tex_idx = tex_idx, tex_size = tex_size, return_type = 'tex')
+                    tex = self.model_forward_synthesis(self.data, self.slic, tex_idx = tex_idx, tex_size = tex_size, return_type = 'tex')
                     col1, col2, col3, col4 = st.columns([1, 1, 4, 1])
                     with col1:
                         st.markdown("")
@@ -140,33 +140,159 @@ class Tester(TesterBase):
                     with col4:
                         st.markdown("")
             st.markdown('<p class="big-font">You can choose another image from the examplar images on the top and start again!</p>', unsafe_allow_html=True)
-            #torch.cuda.empty_cache()
 
-        """
-        st.markdown("#### Texture Editing")
-        st.markdown("**Choose one texture segment to remove.**")
-        remove_idx = clickable_images(
+    def model_forward_synthesis(self, rgb_img, slic, epoch = 1000, test_time = False,
+                                test = True, tex_idx = None, tex_size = 256,
+                                return_type = 'tex', fill_idx = None, remove_idx = None):
+        args = self.args
+        B, _, imgH, imgW = rgb_img.shape
+
+        # Encoder: img (B, 3, H, W) -> feature (B, C, imgH//8, imgW//8)
+        conv_feat, _ = self.model.enc(rgb_img)
+        B, C, H, W = conv_feat.shape
+
+        # Texture code for each superpixel
+        tex_code = self.model.ToTexCode(conv_feat)
+
+        code = F.interpolate(tex_code, size = (imgH, imgW), mode = 'bilinear', align_corners = False)
+        pool_code = poolfeat(code, slic, avg = True)
+
+        prop_code, sp_assign, conv_feats = self.model.gcn(pool_code, slic, (args.add_clustering_epoch <= epoch))
+        softmax = F.softmax(sp_assign * args.temperature, dim = 1)
+        if return_type == 'grouping':
+            return torch.argmax(sp_assign.cpu(), dim = 1)
+
+        tex_seg = poolfeat(conv_feats, softmax, avg = True)
+        seg = label2one_hot_torch(torch.argmax(softmax, dim = 1).unsqueeze(1), C = softmax.shape[1])
+
+        sampled_code = tex_seg[:, tex_idx, :]
+        rec_tex = sampled_code.view(1, -1, 1, 1).repeat(1, 1, tex_size, tex_size)
+        sine_wave = self.model.get_sine_wave(rec_tex, 'rec')
+        H = tex_size // 8; W = tex_size // 8
+        noise = torch.randn(B, self.model.sine_wave_dim, H, W).to(tex_code.device)
+        dec_input = torch.cat((sine_wave, noise), dim = 1)
+
+        weight = self.model.ChannelWeight(rec_tex)
+        weight = F.adaptive_avg_pool2d(weight, output_size = (1)).view(weight.shape[0], -1, 1, 1)
+        weight = torch.sigmoid(weight)
+        dec_input *= weight
+
+        rep_rec = self.model.G(dec_input, rec_tex)
+        rep_rec = (rep_rec + 1) / 2.0
+        return rep_rec
+
+    def display_editing(self):
+        with st.spinner('Running...'):
+            with torch.no_grad():
+                grouping_mask = self.model_forward_editing(self.data, self.slic, return_type = 'grouping')
+
+            data = (self.data + 1) / 2.0
+
+            seg = grouping_mask.view(-1, 1, args.crop_size, args.crop_size)
+            color_vq = self.draw_color_seg(seg)
+            color_vq = color_vq * 0.8 + data.cpu() * 0.2
+
+            st.markdown('<p class="big-font">Given the image you chose, our model decomposes the image into ten texture segments, each depicts one kind of texture in the image.</p>', unsafe_allow_html=True)
+            col1, col2, col3, col4 = st.columns(4)
+            with col1:
+                st.markdown("")
+
+            with col2:
+                st.markdown("Chosen image")
+                st.image(self.to_pil(data))
+
+            with col3:
+                st.markdown("Grouping mask")
+                st.image(self.to_pil(color_vq))
+
+            with col4:
+                st.markdown("")
+
+        seg_onehot = label2one_hot_torch(seg, C = 10)
+        parts = data.cpu() * seg_onehot.squeeze().unsqueeze(1)
+
+        st.markdown('<p class="big-font">We show all texture segments below.</p>', unsafe_allow_html=True)
+        tmp_img_list = []
+        for i in range(parts.shape[0]):
+            part_img = self.to_pil(parts[i])
+            out_path = 'tmp/{}.png'.format(i)
+            part_img.save(out_path)
+
+            with open(out_path, "rb") as image:
+                encoded = base64.b64encode(image.read()).decode()
+                tmp_img_list.append(f"data:image/jpeg;base64,{encoded}")
+
+        tex_idx = clickable_images(
             tmp_img_list,
             titles=[f"Group #{str(i)}" for i in range(len(tmp_img_list))],
             div_style={"display": "flex", "justify-content": "center", "flex-wrap": "wrap"},
-            img_style={"margin": "5px", "height": "120px"},
-            key=1
+            img_style={"margin": "5px", "height": "150px"},
+            key=2
         )
-        st.markdown("**Choose one texture segment to fill in the missing pixels.**")
-        fill_idx = clickable_images(
-            tmp_img_list,
-            titles=[f"Group #{str(i)}" for i in range(len(tmp_img_list))],
+
+        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 = np.asarray(given_mask)
+        given_mask = torch.from_numpy(given_mask)
+        H, W = given_mask.shape[0], given_mask.shape[1]
+        given_mask = label2one_hot_torch(given_mask.view(1, 1, H, W), C = (given_mask.max()+1))
+        mask_img_list = []
+        for i in range(given_mask.shape[1]):
+            part_img = self.to_pil(given_mask[0, i])
+            out_path = 'tmp/{}.png'.format(i)
+            part_img.save(out_path)
+
+            with open(out_path, "rb") as image:
+                encoded = base64.b64encode(image.read()).decode()
+                mask_img_list.append(f"data:image/jpeg;base64,{encoded}")
+
+        part_idx = clickable_images(
+            mask_img_list,
             div_style={"display": "flex", "justify-content": "center", "flex-wrap": "wrap"},
-            img_style={"margin": "5px", "height": "120px"},
-            key=2
+            img_style={"margin": "5px", "height": "150px"},
+            key=1
         )
-        rec = self.model_forward(self.data, self.slic, return_type = 'editing', fill_idx = fill_idx, remove_idx = remove_idx)
-        st.image(self.to_pil(rec))
-        """
 
-    def model_forward(self, rgb_img, slic, epoch = 1000, test_time = False,
+        cols = st.columns(len(mask_img_list))
+        options = []
+        for i, col in enumerate(cols):
+            with col:
+                option = st.selectbox(
+                        "",
+                        ([str(ii) for ii in range(10)]),
+                        key = i)
+                options.append(int(option))
+        print(options)
+
+        if len(options) > 0:
+            with st.spinner('Running...'):
+                st.markdown('<p class="big-font">Edited image is shown below.</p>', unsafe_allow_html=True)
+                #tex_size = st.slider('', 0, 1000, 256)
+                #tex_size = (tex_size // 8) * 8
+                with torch.no_grad():
+                    edited = self.model_forward_editing(self.data, self.slic, options=options, given_mask=given_mask, return_type = 'edited')
+                    col1, col2, col3, col4 = st.columns([1, 1, 4, 1])
+                    with col1:
+                        st.markdown("")
+
+                    with col2:
+                        st.markdown("Input image")
+                        img = F.interpolate(self.data, size = edited.shape[-2:], mode = 'bilinear', align_corners = False)
+                        st.image(self.to_pil((img + 1) / 2.0))
+                        print(img.shape, edited.shape)
+
+                    with col3:
+                        st.markdown("Synthesized texture image")
+                        st.image(self.to_pil(edited))
+
+                    with col4:
+                        st.markdown("")
+            st.markdown('<p class="big-font">You can choose another image from the examplar images on the top and start again!</p>', unsafe_allow_html=True)
+
+    def model_forward_editing(self, rgb_img, slic, epoch = 1000, test_time = False,
                       test = True, tex_idx = None, tex_size = 256,
-                      return_type = 'tex', fill_idx = None, remove_idx = None):
+                      return_type = 'edited', fill_idx = None, remove_idx = None,
+                      options = None, given_mask = None):
         args = self.args
         B, _, imgH, imgW = rgb_img.shape
 
@@ -185,47 +311,29 @@ class Tester(TesterBase):
         if return_type == 'grouping':
             return torch.argmax(sp_assign.cpu(), dim = 1)
 
-
         tex_seg = poolfeat(conv_feats, softmax, avg = True)
         seg = label2one_hot_torch(torch.argmax(softmax, dim = 1).unsqueeze(1), C = softmax.shape[1])
 
-        if return_type == 'tex':
-            sampled_code = tex_seg[:, tex_idx, :]
-            rec_tex = sampled_code.view(1, -1, 1, 1).repeat(1, 1, tex_size, tex_size)
-            sine_wave = self.model.get_sine_wave(rec_tex, 'rec')
-            H = tex_size // 8; W = tex_size // 8
-            noise = torch.randn(B, self.model.sine_wave_dim, H, W).to(tex_code.device)
-            dec_input = torch.cat((sine_wave, noise), dim = 1)
-
-            weight = self.model.ChannelWeight(rec_tex)
-            weight = F.adaptive_avg_pool2d(weight, output_size = (1)).view(weight.shape[0], -1, 1, 1)
-            weight = torch.sigmoid(weight)
-            dec_input *= weight
-
-            rep_rec = self.model.G(dec_input, rec_tex)
-            rep_rec = (rep_rec + 1) / 2.0
-            return rep_rec
-        elif return_type == 'editing':
-            remove_mask = 0
-            fill_mask = 1
-            rec_tex = upfeat(tex_seg, seg)
-            remove_mask = seg[:, remove_idx:remove_idx+1]
-            fill_tex = tex_seg[:, fill_idx, :].view(1, -1, 1, 1).repeat(1, 1, imgH, imgW)
-            rec_tex = rec_tex * (1 - remove_mask) + fill_tex * remove_mask
-
-            sine_wave = self.model.get_sine_wave(rec_tex, 'rec')
-            H = imgH // 8; W = imgW // 8
-            noise = torch.randn(B, self.model.sine_wave_dim, H, W).to(tex_code.device)
-            dec_input = torch.cat((sine_wave, noise), dim = 1)
-            weight = self.model.ChannelWeight(rec_tex)
-            weight = F.adaptive_avg_pool2d(weight, output_size = (1)).view(weight.shape[0], -1, 1, 1)
-            weight = torch.sigmoid(weight)
-            dec_input *= weight
-
-            rep_rec = self.model.G(dec_input, rec_tex)
-            rep_rec = (rep_rec + 1) / 2.0
-            return rep_rec
-
+        given_mask = F.interpolate(given_mask, size = (512, 512), mode = 'bilinear', align_corners = False)
+        rec_tex = torch.zeros((1, tex_seg.shape[-1], 512, 512))
+        for i in range(given_mask.shape[1]):
+            label = options[i]
+            code = tex_seg[0, label, :].view(1, -1, 1, 1).repeat(1, 1, 512, 512)
+            rec_tex += code * given_mask[:, i:i+1]
+        tex_size = 512
+        sine_wave = self.model.get_sine_wave(rec_tex, 'rec')
+        H = tex_size // 8; W = tex_size // 8
+        noise = torch.randn(B, self.model.sine_wave_dim, H, W).to(tex_code.device)
+        dec_input = torch.cat((sine_wave, noise), dim = 1)
+
+        weight = self.model.ChannelWeight(rec_tex)
+        weight = F.adaptive_avg_pool2d(weight, output_size = (1)).view(weight.shape[0], -1, 1, 1)
+        weight = torch.sigmoid(weight)
+        dec_input *= weight
+
+        rep_rec = self.model.G(dec_input, rec_tex)
+        rep_rec = (rep_rec + 1) / 2.0
+        return rep_rec
 
     def load_data(self, data_path):
         rgb_img = Image.open(data_path)
@@ -253,12 +361,12 @@ class Tester(TesterBase):
         self.model = self.model.module
         return
 
+    """
     def test(self):
-        """ Test function
-        """
         #for iteration in tqdm(range(args.nsamples)):
         self.test_step(0)
         self.display(0, 'train')
+    """
 
 def main():
     #torch.cuda.empty_cache()
@@ -300,7 +408,13 @@ def main():
         tester.define_model()
         tester.load_data(img_path)
         tester.load_model(args.pretrained_path)
-        tester.display()
+        tab1, tab2 = st.tabs(["Texture Synthesis", "Texture Editing"])
+        with tab1:
+            st.header("Texture Synthesis")
+            tester.display_synthesis()
+        with tab2:
+            st.header("Texture Editing")
+            tester.display_editing()
 
 if __name__ == '__main__':
     os.system("pip install torch-geometric==1.7.2")