Modified vispy's turntable camera

newttcam.py

# -*- coding: utf-8 -*-
# Copyright (c) Vispy Development Team. All Rights Reserved.
# Distributed under the (new) BSD License. See LICENSE.txt for more info.

from __future__ import division

import math
import numpy as np

from vispy.scene.cameras.base_camera import BaseCamera
from vispy.util import keys, transforms
from vispy.visuals.transforms import MatrixTransform
from vispy.scene.cameras import PerspectiveCamera


def as_vec4(obj, default=(0, 0, 0, 1)):
    """
    Convert `obj` to 4-element vector (numpy array with shape[-1] == 4)

    Parameters
    ----------
    obj : array-like
        Original object.
    default : array-like
        The defaults to use if the object does not have 4 entries.

    Returns
    -------
    obj : array-like
        The object promoted to have 4 elements.

    Notes
    -----
    `obj` will have at least two dimensions.

    If `obj` has < 4 elements, then new elements are added from `default`.
    For inputs intended as a position or translation, use default=(0,0,0,1).
    For inputs intended as scale factors, use default=(1,1,1,1).

    """
    obj = np.atleast_2d(obj)
    # For multiple vectors, reshape to (..., 4)
    if obj.shape[-1] < 4:
        new = np.empty(obj.shape[:-1] + (4,), dtype=obj.dtype)
        new[:] = default
        new[..., :obj.shape[-1]] = obj
        obj = new
    elif obj.shape[-1] > 4:
        raise TypeError("Array shape %s cannot be converted to vec4"
                        % (obj.shape, ))
    return obj


class NewTurntableCamera(PerspectiveCamera):
    _state_props = PerspectiveCamera._state_props + ('elevation',
                                                        'azimuth', 'roll')

    def __init__(self, fov=45.0, elevation=30.0, azimuth=30.0, roll=0.0,
                 distance=None, translate_speed=1.0, **kwargs):
        self.__init2__(fov=fov, **kwargs)

        # Set camera attributes
        self.azimuth = azimuth
        self.elevation = elevation
        self.roll = roll  # interaction not implemented yet
        self.distance = distance  # None means auto-distance
        self.translate_speed = translate_speed

    @property
    def elevation(self):
        """ The angle of the camera in degrees above the horizontal (x, z)
        plane.
        """
        return self._elevation

    @elevation.setter
    def elevation(self, elev):
        elev = float(elev)
        self._elevation = min(90, max(-90, elev))
        self.view_changed()

    @property
    def azimuth(self):
        """ The angle of the camera in degrees around the y axis. An angle of
        0 places the camera within the (y, z) plane.
        """
        return self._azimuth

    @azimuth.setter
    def azimuth(self, azim):
        azim = float(azim)
        while azim < -180:
            azim += 360
        while azim > 180:
            azim -= 360
        self._azimuth = azim
        self.view_changed()

    @property
    def roll(self):
        """ The angle of the camera in degrees around the z axis. An angle of
        0 places puts the camera upright.
        """
        return self._roll

    @roll.setter
    def roll(self, roll):
        roll = float(roll)
        while roll < -180:
            roll += 360
        while roll > 180:
            roll -= 360
        self._roll = roll
        self.view_changed()

    def orbit(self, azim, elev):
        """ Orbits the camera around the center position.

        Parameters
        ----------
        azim : float
            Angle in degrees to rotate horizontally around the center point.
        elev : float
            Angle in degrees to rotate vertically around the center point.
        """
        self.azimuth += azim
        self.elevation = np.clip(self.elevation + elev, -90, 90)
        self.view_changed()

    def _update_rotation(self, event):
        """Update rotation parmeters based on mouse movement"""
        p1 = event.mouse_event.press_event.pos
        p2 = event.mouse_event.pos
        if self._event_value is None:
            self._event_value = self.azimuth, self.elevation
        self.azimuth = self._event_value[0] - (p2 - p1)[0] * 0.5
        self.elevation = self._event_value[1] + (p2 - p1)[1] * 0.5


    def _update_camera_pos(self):
        """ Set the camera position and orientation"""

        # transform will be updated several times; do not update camera
        # transform until we are done.
        ch_em = self.events.transform_change
        with ch_em.blocker(self._update_transform):
            tr = self.transform

            up, forward, right = self._get_dim_vectors()

            # Create mapping so correct dim is up
            pp1 = np.array([(0, 0, 0), (0, 0, -1), (1, 0, 0), (0, 1, 0)])
            pp2 = np.array([(0, 0, 0), forward, right, up])
            
            # tr.set_mapping(pp1, pp2)
            m = transforms.affine_map(pp1, pp2).T

            # tr.translate(-self._actual_distance * forward)
            pos = -self._actual_distance * forward
            m = np.dot(m, transforms.translate(pos))

            """Rotate the transformation matrix based on camera parameters"""
            up, forward, right = self._get_dim_vectors()
            # tr.rotate(self.elevation, -right)
            m = np.dot(m, transforms.rotate(self.elevation, -right))
            # tr.rotate(self.azimuth, up)
            m = np.dot(m, transforms.rotate(self.azimuth, up))
            # tr.scale([1.0/a for a in self._flip_factors])
            scale = [1.0/a for a in self._flip_factors]
            m = np.dot(m, transforms.scale(as_vec4(scale, default=(1, 1, 1, 1))[0, :3]))
            # tr.translate(np.array(self.center))
            pos = np.array(self.center)
            m = np.dot(m, transforms.translate(pos))

            tr.matrix = m

    def _dist_to_trans(self, dist):
        """Convert mouse x, y movement into x, y, z translations"""
        rae = np.array([self.roll, self.azimuth, self.elevation]) * np.pi / 180
        sro, saz, sel = np.sin(rae)
        cro, caz, cel = np.cos(rae)
        d0, d1 = dist[0], dist[1]
        dx = (+ d0 * (cro * caz + sro * sel * saz)
              + d1 * (sro * caz - cro * sel * saz)) * self.translate_speed
        dy = (+ d0 * (cro * saz - sro * sel * caz)
              + d1 * (sro * saz + cro * sel * caz)) * self.translate_speed
        dz = (- d0 * sro * cel + d1 * cro * cel) * self.translate_speed
        return dx, dy, dz

    def __init2__(self, fov=0.0, **kwargs):
        super(NewTurntableCamera, self).__init__(fov=fov, **kwargs)
        self._actual_distance = 0.0
        self._event_value = None

    @property
    def distance(self):
        """ The user-set distance. If None (default), the distance is
        internally calculated from the scale factor and fov.
        """
        return self._distance

    @distance.setter
    def distance(self, distance):
        if distance is None:
            self._distance = None
        else:
            self._distance = float(distance)
        self.view_changed()

    def viewbox_mouse_event(self, event):
        """
        The viewbox received a mouse event; update transform
        accordingly.

        Parameters
        ----------
        event : instance of Event
            The event.
        """
        if event.handled or not self.interactive:
            return

        PerspectiveCamera.viewbox_mouse_event(self, event)

        if event.type == 'mouse_release':
            self._event_value = None  # Reset
        elif event.type == 'mouse_press':
            event.handled = True
        elif event.type == 'mouse_move':
            if event.press_event is None:
                return

            modifiers = event.mouse_event.modifiers
            p1 = event.mouse_event.press_event.pos
            p2 = event.mouse_event.pos
            d = p2 - p1

            if 1 in event.buttons and not modifiers:
                # Rotate
                self._update_rotation(event)

            elif 2 in event.buttons and not modifiers:
                # Zoom
                if self._event_value is None:
                    self._event_value = (self._scale_factor, self._distance)
                zoomy = (1 + self.zoom_factor) ** d[1]
                
                self.scale_factor = self._event_value[0] * zoomy
                # Modify distance if its given
                if self._distance is not None:
                    self._distance = self._event_value[1] * zoomy
                self.view_changed()

            elif 1 in event.buttons and keys.SHIFT in modifiers:
                # Translate
                norm = np.mean(self._viewbox.size)
                if self._event_value is None or len(self._event_value) == 2:
                    self._event_value = self.center
                dist = (p1 - p2) / norm * self._scale_factor
                dist[1] *= -1
                # Black magic part 1: turn 2D into 3D translations
                dx, dy, dz = self._dist_to_trans(dist)
                # Black magic part 2: take up-vector and flipping into account
                ff = self._flip_factors
                up, forward, right = self._get_dim_vectors()
                dx, dy, dz = right * dx + forward * dy + up * dz
                dx, dy, dz = ff[0] * dx, ff[1] * dy, dz * ff[2]
                c = self._event_value
                self.center = c[0] + dx, c[1] + dy, c[2] + dz

            elif 2 in event.buttons and keys.SHIFT in modifiers:
                # Change fov
                if self._event_value is None:
                    self._event_value = self._fov
                fov = self._event_value - d[1] / 5.0
                self.fov = min(180.0, max(0.0, fov))

    def _get_dim_vectors(self):
        # Specify up and forward vector
        M = {'+z': [(0, 0, +1), (0, 1, 0)],
             '-z': [(0, 0, -1), (0, 1, 0)],
             '+y': [(0, +1, 0), (1, 0, 0)],
             '-y': [(0, -1, 0), (1, 0, 0)],
             '+x': [(+1, 0, 0), (0, 0, 1)],
             '-x': [(-1, 0, 0), (0, 0, 1)],
             }
        up, forward = M[self.up]
        right = np.cross(forward, up)
        return np.array(up), np.array(forward), right

    def _update_projection_transform(self, fx, fy):
        d = self.depth_value
        if self._fov == 0:
            self._projection.set_ortho(-0.5*fx, 0.5*fx, -0.5*fy, 0.5*fy, -d, d)
            self._actual_distance = self._distance or 0.0
        else:
            # Figure distance to center in order to have correct FoV and fy.
            # Use that auto-distance, or the given distance (if not None).
            fov = max(0.01, self._fov)
            dist = fy / (2 * math.tan(math.radians(fov)/2))
            self._actual_distance = dist = self._distance or dist
            val = math.sqrt(d*10)
            self._projection.set_perspective(fov, fx/fy, dist/val, dist*val)
        # Update camera pos, which will use our calculated _distance to offset
        # the camera
        self._update_camera_pos()