Merge pull request #65 from PDillis/main

General fixes

visualizer_drag.py

@@ -177,15 +177,16 @@ class Visualizer(imgui_window.ImguiWindow):
                     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):

viz/capture_widget.py

@@ -31,7 +31,7 @@ class CaptureWidget:
         viz = self.viz
         try:
             _height, _width, channels = image.shape
-            assert channels in [1, 3]
+            print(viz.result)
             assert image.dtype == np.uint8
             os.makedirs(self.path, exist_ok=True)
             file_id = 0
@@ -43,8 +43,9 @@ class CaptureWidget:
             if channels == 1:
                 pil_image = PIL.Image.fromarray(image[:, :, 0], 'L')
             else:
-                pil_image = PIL.Image.fromarray(image, 'RGB')
+                pil_image = PIL.Image.fromarray(image[:, :, :3], 'RGB')
             pil_image.save(os.path.join(self.path, f'{file_id:05d}.png'))
+            np.save(os.path.join(self.path, f'{file_id:05d}.npy'), viz.result.w)
         except:
             viz.result.error = renderer.CapturedException()
 

viz/renderer.py

@@ -382,5 +382,6 @@ class Renderer:
             img = img.cpu().numpy()
             img = Image.fromarray(img)
         res.image = img
+        res.w = ws.detach().cpu().numpy()
 
 #----------------------------------------------------------------------------