Actual fixes for resetting mask and saving w.

PeacefulHorizon · Jun 26, 2023 · 28a3e17 · 28a3e17
1 parent b76c49a
commit 28a3e17
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/visualizer_drag.py b/visualizer_drag.py
@@ -177,15 +177,16 @@ def draw_frame(self):
                     if self.result.init_net:
                         self.drag_widget.reset_point()
 
-        # if self.check_update_mask(**self.args):
-        #     h, w, _ = self.result.image.shape
-        #     self.drag_widget.init_mask(w, h)
-
         # Display.
         max_w = self.content_width - self.pane_w
         max_h = self.content_height
         pos = np.array([self.pane_w + max_w / 2, max_h / 2])
         if 'image' in self.result:
+            # Reset mask after loading a new pickle or changing seed.
+            if self.check_update_mask(**self.args):
+                h, w, _ = self.result.image.shape
+                self.drag_widget.init_mask(w, h)
+
             if self._tex_img is not self.result.image:
                 self._tex_img = self.result.image
                 if self._tex_obj is None or not self._tex_obj.is_compatible(image=self._tex_img):

diff --git a/viz/renderer.py b/viz/renderer.py
@@ -382,6 +382,6 @@ def _render_drag_impl(self, res,
             img = img.cpu().numpy()
             img = Image.fromarray(img)
         res.image = img
-        res.w = ws.cpu().numpy()
+        res.w = ws.detach().cpu().numpy()
 
 #----------------------------------------------------------------------------