Merge pull request #44 from aylward/scene_screenshot

ENH: Save a screenshot for 3D view when saving a scene

src/minder3d/lib/sovImageTablePanelWidget.py

@@ -122,7 +122,7 @@ def unload_selected(self):
                 self.gui.unload_scene(False)
         if len(self.selected) > 0:
             self.selected = []
-            self.fill_table()
+        self.fill_table()
 
     @time_and_log
     def remove_selected(self):
@@ -146,7 +146,7 @@ def remove_selected(self):
                 )
         if len(self.selected) > 0:
             self.selected = []
-            self.fill_table()
+        self.fill_table()
 
     @time_and_log
     def close_expanded_table(self):
@@ -475,7 +475,9 @@ def save_image(self, filename):
     @time_and_log
     def load_scene(self):
         self.state.scene_thumbnail = self.settings.get_thumbnail(
-            self.state.scene, self.state.scene_filename, 'scene'
+            self.gui.view3DPanel.get_screenshot(),
+            self.state.scene_filename,
+            'scene',
         )
         self.state.scene_label = os.path.basename(self.state.scene_filename)
         self.settings.add_data(
@@ -490,13 +492,18 @@ def load_scene(self):
 
     @time_and_log
     def save_scene(self, filename):
+        self.state.scene_thumbnail = self.settings.get_thumbnail(
+            self.gui.view3DPanel.get_screenshot(), filename, 'scene', True
+        )
+        self.state.scene_label = os.path.basename(filename)
         self.settings.add_data(
             self.state.scene,
             filename,
             'scene',
             self.state.scene_label,
             self.state.scene_thumbnail,
         )
+        self.fill_table()
 
     @time_and_log
     def replace_image(self, img_num):
@@ -561,7 +568,7 @@ def load_selected(self):
         self.fill_table()
 
     @time_and_log
-    def register_images(self, dir, first=True):
+    def register_images(self, dir, redraw_table=True):
         files = [
             os.path.join(dir, f)
             for f in os.listdir(dir)
@@ -584,7 +591,6 @@ def register_images(self, dir, first=True):
                 continue
             label = os.path.basename(filename)
             thumbnail = self.settings.get_thumbnail(img, filename, 'image')
-            print(f'Adding {filename} to settings')
             self.settings.add_data(
                 img,
                 filename,
@@ -601,6 +607,6 @@ def register_images(self, dir, first=True):
         for dirname in dirs:
             self.register_images(dirname, False)
 
-        if first:
+        if redraw_table:
             self.selected = []
             self.fill_table()

src/minder3d/lib/sovImageTableSettings.py

@@ -5,6 +5,7 @@
 import numpy as np
 from PySide6.QtCore import QSettings, QStandardPaths, Qt
 from PySide6.QtGui import QImage, QPixmap
+from vtk.util.numpy_support import vtk_to_numpy
 
 from .sovUtils import time_and_log
 
@@ -217,18 +218,22 @@ def relabel(self, filename, new_label):
                 return
 
     @time_and_log
-    def get_thumbnail(self, obj, filename, file_type):
+    def get_thumbnail(
+        self, obj, filename, file_type, force_new_thumbnail=False
+    ):
         """Get the thumbnail of a file.
 
         This function retrieves the thumbnail of a file from the settings.
         """
-        for file in self.file_records:
-            if file.filename == filename:
-                return QPixmap(file.file_thumbnail)
+        if not force_new_thumbnail:
+            for file in self.file_records:
+                if file.filename == filename:
+                    return QPixmap(file.file_thumbnail)
+
         if file_type == 'image':
             return self.get_thumbnail_pixmap_from_image(obj)
         elif file_type == 'scene':
-            return QPixmap(':/icons/scene.png')
+            return self.get_thumbnail_pixmap_from_vtk_image(obj)
 
     @time_and_log
     def get_thumbnail_pixmap_from_image(self, img):
@@ -259,3 +264,40 @@ def get_thumbnail_pixmap_from_image(self, img):
             100, 100, Qt.KeepAspectRatio
         )
         return thumb_pixmap
+
+    @time_and_log
+    def get_thumbnail_pixmap_from_vtk_image(self, img):
+        vtk_array = img.GetPointData().GetScalars()
+        np_array = vtk_to_numpy(vtk_array)
+        dims = img.GetDimensions()
+        np_array = np_array.reshape(dims[1], dims[0], dims[2], -1)
+
+        # Process the NumPy array
+        if np_array.shape[3] == 1:
+            np_array = np_array[:, :, :, 0]
+            np_array = np.stack([np_array] * 3, axis=-1)
+        elif np_array.shape[3] == 4:
+            np_array = np_array[:, :, :, :3]
+
+        auto_range = np.quantile(np_array, [0.05, 0.95])
+        np_array = np.clip(np_array, auto_range[0], auto_range[1])
+        np_array = (
+            (np_array - auto_range[0]) / (auto_range[1] - auto_range[0]) * 255
+        ).astype(np.uint8)
+
+        # Create QImage
+        thumb_image = QImage(
+            np_array.data,
+            np_array.shape[1],
+            np_array.shape[0],
+            np_array.strides[0],
+            QImage.Format_RGB888,
+        )
+        thumb_image.setDotsPerMeterX(10 / img.GetSpacing()[0])
+        thumb_image.setDotsPerMeterY(10 / img.GetSpacing()[1])
+
+        # Convert QImage to QPixmap
+        thumb_pixmap = QPixmap.fromImage(thumb_image).scaled(
+            100, 100, Qt.KeepAspectRatio
+        )
+        return thumb_pixmap

src/minder3d/lib/sovView3DPanelWidget.py

@@ -1,7 +1,7 @@
-import numpy as np
 from PySide6.QtWidgets import QWidget
-from vtk import vtkRenderLargeImage
-from vtk.util.numpy_support import vtk_to_numpy
+
+# from vtk import vtkRenderLargeImage
+from vtk import vtkWindowToImageFilter
 
 from .sovUtils import time_and_log
 from .sovView3DRenderWindowInteractor import View3DRenderWindowInteractor
@@ -35,6 +35,7 @@ def reset_camera(self):
         self.vtk3DViewWidget.reset_camera()
 
     def get_screenshot(self):
+        """
         render = vtkRenderLargeImage()
         render.SetMagnification(1)
         render.SetInput(self.vtk3DViewWidget.scene_renderer)
@@ -49,6 +50,17 @@ def get_screenshot(self):
         numpy_data = numpy_data.transpose(0, 1, 2)
         numpy_data = np.flipud(numpy_data)
         return numpy_data
+        """
+        screenshotfilter = vtkWindowToImageFilter()
+        screenshotfilter.SetInput(
+            self.vtk3DViewWidget.scene_renderer.GetRenderWindow()
+        )
+        screenshotfilter.SetScale(2)
+        screenshotfilter.SetInputBufferTypeToRGBA()
+        screenshotfilter.ReadFrontBufferOff()
+        screenshotfilter.Update()
+
+        return screenshotfilter.GetOutput()
 
     @time_and_log
     def create_new_image(self):

src/minder3d/minder3DWindow.py

@@ -347,7 +347,9 @@ def save_vtk_models(self, filename=None):
             )
         if filename:
             exporter = vtk.vtkVRMLExporter()
-            exporter.SetRenderWindow(self.view3DPanel.renderWindow)
+            exporter.SetRenderWindow(
+                self.view3DPanel.vtk3DViewWidget.GetRenderWindow()
+            )
             exporter.SetFileName(filename)
             exporter.Write()
             exporter.Update()
@@ -374,7 +376,6 @@ def save_scene(self, filename=None):
             self.state.scene_filename = os.path.abspath(filename)
             self.log(f'Saving scene to {filename}')
             new_scene = compress_scene_for_saving(self.state.scene)
-            print(new_scene)
             write_group(new_scene, filename)
             self.imageTablePanel.save_scene(filename)