add 3D dendrogram selection but not got dendrogram data

glue_vispy_viewers/common/toolbar.py

@@ -24,6 +24,7 @@
 ROTATE_ICON = os.path.join(os.path.dirname(__file__), 'glue_rotate.png')
 RECORD_START_ICON = os.path.join(os.path.dirname(__file__), 'glue_record_start.png')
 RECORD_STOP_ICON = os.path.join(os.path.dirname(__file__), 'glue_record_stop.png')
+DENDROGRAM_ICON = os.path.join(os.path.dirname(__file__), 'glue_dendrogram.png')
 
 
 class PatchedElementSubsetState(ElementSubsetState):
@@ -135,6 +136,14 @@ def __init__(self, vispy_widget=None, parent=None):
         self.addAction(a)
         self.point_action = a
 
+        # TODO: show if there is dendrogram file or not shown
+        a = QtGui.QAction(QtGui.QIcon(DENDROGRAM_ICON), 'Dendrogram Selection', parent)
+        a.triggered.connect(nonpartial(self.toggle_dendrogram))
+        a.setCheckable(True)
+        parent.addAction(a)
+        self.addAction(a)
+        self.dendrogram_action = a
+
         # Connect callback functions to VisPy Canvas
         self._vispy_widget.canvas.events.mouse_press.connect(self.on_mouse_press)
         self._vispy_widget.canvas.events.mouse_release.connect(self.on_mouse_release)
@@ -185,43 +194,51 @@ def record(self, event):
     def toggle_lasso(self):
         if self.lasso_action.isChecked():
             self.mode = 'lasso'
-            self.rectangle_action.setChecked(False)
-            self.point_action.setChecked(False)
-            self.ellipse_action.setChecked(False)
-            self.rotate_action.setChecked(False)
+            self.set_all_false()
+            self.lasso_action.setChecked(True)
         else:
             self.mode = None
 
     def toggle_rectangle(self):
         if self.rectangle_action.isChecked():
             self.mode = 'rectangle'
-            self.lasso_action.setChecked(False)
-            self.point_action.setChecked(False)
-            self.ellipse_action.setChecked(False)
-            self.rotate_action.setChecked(False)
+            self.set_all_false()
+            self.rectangle_action.setChecked(True)
         else:
             self.mode = None
 
     def toggle_ellipse(self):
         if self.ellipse_action.isChecked():
             self.mode = 'ellipse'
-            self.lasso_action.setChecked(False)
-            self.point_action.setChecked(False)
-            self.rectangle_action.setChecked(False)
-            self.rotate_action.setChecked(False)
+            self.set_all_false()
+            self.ellipse_action.setChecked(True)
         else:
             self.mode = None
 
     def toggle_point(self):
         if self.point_action.isChecked():
             self.mode = 'point'
-            self.lasso_action.setChecked(False)
-            self.rectangle_action.setChecked(False)
-            self.ellipse_action.setChecked(False)
-            self.rotate_action.setChecked(False)
+            self.set_all_false()
+            self.point_action.setChecked(True)
         else:
             self.mode = None
 
+    def toggle_dendrogram(self):
+        if self.dendrogram_action.isChecked():
+            self.mode = 'dendrogram'
+            self.set_all_false()
+            self.dendrogram_action.setChecked(True)
+        else:
+            self.mode = None
+
+    def set_all_false(self):
+        self.lasso_action.setChecked(False)
+        self.rectangle_action.setChecked(False)
+        self.ellipse_action.setChecked(False)
+        self.point_action.setChecked(False)
+        self.dendrogram_action.setChecked(False)
+        self.rotate_action.setChecked(False)
+
     def toggle_rotate(self):
         if self.rotate_action.isChecked():
             # Start the rotation

glue_vispy_viewers/volume/volume_toolbar.py

@@ -11,28 +11,100 @@
 from matplotlib import path
 from glue.core.roi import RectangularROI, CircularROI, PolygonalROI
 from astrodendro import Dendrogram
+from ..extern.vispy import scene
 
 
 class VolumeSelectionToolbar(VispyDataViewerToolbar):
 
     def __init__(self, vispy_widget=None, parent=None):
         super(VolumeSelectionToolbar, self).__init__(vispy_widget=vispy_widget, parent=parent)
+        # for getting transposed vol_data shape and getting pos array
         self.trans_ones_data = None
+        # add some markers
+        self.markers = scene.visuals.Markers(parent=self._vispy_widget.view.scene)
+        self.ray_line = scene.visuals.Line(color='green', width=5,
+                                           parent=self._vispy_widget.view.scene)
+
+        self.visual_tr = None
+        self.visible_data = None
+        self.visual = None
+        self.vol_data = None
+        self.max_value_pos = None
 
     def on_mouse_press(self, event):
-        self.selection_origin = event.pos
-        if self.mode is 'point':
-            visible_data, visual = self.get_visible_data()
-            current_layer = visible_data[0]
-            max_pos = self.get_max_pos()
+        """
+        Assign mouse position and do point selection.
+        :param event:
+        """
+        if self.mode:
+            # do the initiation here
+            self.selection_origin = event.pos
+            self.visible_data, self.visual = self.get_visible_data()
+
+            # for dendrogram cube the two attributes are intensity and structure, no ['PRIMARY']
+            data_array = self.visible_data[0]['intensity']
+
+            self.trans_ones_data = np.transpose(np.ones(data_array.shape))
+
+            self.vol_data = data_array # current_layer
+            self.visual_tr = self._vispy_widget.limit_transforms[self.visual]
+
+        if self.mode is 'dendrogram':
+
+            # get start and end point of ray line
+            pos = self.get_ray_line()
+            max_value_pos, max_value = self.get_inter_value(pos)
+            self.max_value_pos = max_value_pos
+
+            # change the color of max value marker
+            if max_value:
+                self.markers.set_data(pos=np.array(max_value_pos),
+                                      face_color='yellow')
+                status_text = 'pos '+str(max_value_pos[0]) \
+                              + ' value '+str(max_value)
+                # TODO: AttributeError: 'VispyVolumeViewer' object has no attribute 'show_status'
+                # self._vispy_data_viewer.show_status(status_text)
+
+                # TODO: vispy_viewer doesn't have client, how to get dendrogram data?
+                print('visible_data, client.data', self.visible_data, self._vispy_widget.client.data)
+                for each_data in self._vispy_widget.client.data:
+                    if each_data.label == 'Dendrogram':
+                        d = each_data
+
+                try:
+                    branch_label = self.visible_data[0]['structure'][max_value_pos]
+                    print('branch_label', branch_label)
+                except IndexError:
+                    branch_label = -1
+
+                # no structure found
+                if branch_label == -1:
+                    return None
+
+                # similar to DFS to find children
+                def get_all_branch(d, branch_label, ini_mask):
+                    mask = np.logical_or(ini_mask, self.visible_data['structure'] == branch_label)
+                    child_label = np.where(d['parent'] == branch_label)
+
+                    if len(child_label[0]) != 0:
+                        for each_child in child_label[0]:
+                            mask = get_all_branch(d, each_child, mask)
+                    return mask
+
+                ini_mask = np.zeros(self.visible_data['structure'].shape, dtype=bool)
+                mask = get_all_branch(d, branch_label, ini_mask)
+                print('sum mask', sum(mask))
+
+                # dendro_mask = substructure.get_mask(shape=np.transpose(self.trans_ones_data).shape)
+                # dendro_mask = np.ravel(dendro_mask) # confused by the transpose here
+                self.mark_selected(mask, self.vol_data)
 
-            dendro = Dendrogram.load_from('dendrogram.fits') # TODO: find the glue dendro structure
-            # maxpos regards to the transposed array while dendrogram to the original data
-            substructure = dendro.structure_at((max_pos[0], max_pos[2], max_pos[1]))
+                self._vispy_widget.canvas.update()
 
-            dendro_mask = substructure.get_mask(shape=np.transpose(self.trans_ones_data).shape)
-            # dendro_mask = np.ravel(dendro_mask) # confused by the transpose here
-            self.mark_selected(dendro_mask, current_layer)
+                # visible_data, visual = self.get_visible_data()
+                # self.vol_data = visible_data
+                # current_layer = visible_data[0]
+                # max_pos = self.get_max_pos()
 
     def on_mouse_release(self, event):
 
@@ -80,18 +152,23 @@ def on_mouse_release(self, event):
             self.lasso_reset()
 
     def get_map_data(self):
+        """
+        Get the mapped buffer from self.visual to canvas.
+        :return: Mapped data position on canvas.
+        """
 
         # Get the visible datasets
-        visible_data, visual = self.get_visible_data()
-        layer = visible_data[0]
+        data_object = self.visible_data[0]
+
+        # TODO: add support for multiple data here, data_array should
+        # cover all self.visible_data array
 
-        # TODO: add support for multiple data here, layer should cover all visible_data array
+        tr = as_matrix_transform(self.visual.get_transform(map_from='visual', map_to='canvas'))
 
-        tr = as_matrix_transform(visual.get_transform(map_from='visual', map_to='canvas'))
+        self.trans_ones_data = np.transpose(np.ones(data_object.data.shape))
 
-        self.trans_ones_data = np.transpose(np.ones(layer.data.shape))
+        pos_data = np.indices(data_object.data.shape[::-1], dtype=float).reshape(3, -1).transpose()
 
-        pos_data = np.indices(layer.data.shape[::-1], dtype=float).reshape(3, -1).transpose()
         data = tr.map(pos_data)
         data /= data[:, 3:]   # normalize with homogeneous coordinates
         return data[:, :2]
@@ -112,3 +189,68 @@ def get_max_pos(self):
         print('maxpos, maxvalue', max_pos, max_value)
         # return max_pos[0], max_pos[1], max_pos[2]
         return max_pos # regarding to the transposed volume data
+
+    def get_inter_value(self, pos):
+        trans = self.visual_tr.translate
+        scale = self.visual_tr.scale
+
+        inter_pos = []
+        for z in range(0, self.vol_data.shape[0]):
+            #   3D line defined with two points (x0, y0, z0) and (x1, y1, z1) as
+            #   (x - x1)/(x2 - x1) = (y - y1)/(y2 - y1) = (z - z1)/(z2 - z1) = t
+            z = z * scale[2] + trans[2]
+            t = (z - pos[0][2])/(pos[1][2] - pos[0][2])
+            x = t * (pos[1][0] - pos[0][0]) + pos[0][0]
+            y = t * (pos[1][1] - pos[0][1]) + pos[0][1]
+            inter_pos.append([x, y, z])
+        inter_pos = np.array(inter_pos)
+
+        # cut the line within the cube
+        m1 = inter_pos[:, 0] > trans[0]  # for x
+        m2 = inter_pos[:, 0] < (self.vol_data.shape[2] * scale[0] + trans[0])
+        m3 = inter_pos[:, 1] > trans[1]  # for y
+        m4 = inter_pos[:, 1] < (self.vol_data.shape[1]*scale[1] + trans[1])
+        inter_pos = inter_pos[m1 & m2 & m3 & m4]
+
+        # set colors for markers
+        colors = np.ones((inter_pos.shape[0], 4))
+        colors[:] = (0.5, 0.5, 0, 1)
+
+        # value of intersected points
+        inter_value = []
+        for each_point in inter_pos:
+            x = (each_point[0] - trans[0])/scale[0]
+            y = (each_point[1] - trans[1])/scale[1]
+            z = (each_point[2] - trans[2])/scale[2]
+            inter_value.append(self.vol_data[(z, y, x)])
+        inter_value = np.array(inter_value)
+
+        assert inter_value.shape[0] == inter_pos.shape[0]
+
+        # TODO: there is a risk that no intersected points found here
+        if len(inter_pos) == 0 or len(inter_value) == 0:
+            print('empty selection list', inter_pos, inter_value)
+            return None, None
+        else:
+            return [inter_pos[np.argmax(inter_value)]], inter_value.max()
+
+    def get_ray_line(self):
+        """
+        Get the ray line from camera pos to the far point.
+        :return: Start point and end point position.
+        """
+        tr_back = self.visual.get_transform(map_from='canvas', map_to='visual')
+
+        _cam = self._vispy_widget.view.camera
+        start_point = _cam.transform.map(_cam.center)
+
+        end_point = np.insert(self.selection_origin, 2, 1)
+        end_point = tr_back.map(end_point)
+
+        # add the self.visual local transform
+        end_point = self.visual_tr.map(end_point)
+        end_point = end_point[:3] / end_point[3]
+
+        # self.ray_line.set_data(np.array([end_point, start_point[:3]]))
+
+        return np.array([end_point, start_point[:3]])