path: root/viewport_vr_preview.py

# ##### BEGIN GPL LICENSE BLOCK #####
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####

# <pep8 compliant>

import bpy
from mathutils import Matrix, Quaternion, Vector
from bpy.types import (
    Gizmo,
    GizmoGroup,
)
from bpy.props import (
    CollectionProperty,
    IntProperty,
    BoolProperty,
)
from bpy.app.handlers import persistent

bl_info = {
    "name": "VR Scene Inspection",
    "author": "Julian Eisel (Severin)",
    "version": (0, 2, 0),
    "blender": (2, 83, 8),
    "location": "3D View > Sidebar > VR",
    "description": ("View the viewport with virtual reality glasses "
                    "(head-mounted displays)"),
    "support": "OFFICIAL",
    "warning": "This is an early, limited preview of in development "
               "VR support for Blender.",
    "doc_url": "{BLENDER_MANUAL_URL}/addons/3d_view/vr_scene_inspection.html",
    "category": "3D View",
}


@persistent
def ensure_default_vr_landmark(context: bpy.context):
    # Ensure there's a default landmark (scene camera by default).
    landmarks = bpy.context.scene.vr_landmarks
    if not landmarks:
        landmarks.add()
        landmarks[0].type = 'SCENE_CAMERA'


def xr_landmark_active_type_update(self, context):
    wm = context.window_manager
    session_settings = wm.xr_session_settings
    landmark_active = VRLandmark.get_active_landmark(context)

    # Update session's base pose type to the matching type.
    if landmark_active.type == 'SCENE_CAMERA':
        session_settings.base_pose_type = 'SCENE_CAMERA'
    elif landmark_active.type == 'USER_CAMERA':
        session_settings.base_pose_type = 'OBJECT'
    # elif landmark_active.type == 'CUSTOM':
        # session_settings.base_pose_type = 'CUSTOM'


def xr_landmark_active_camera_update(self, context):
    session_settings = context.window_manager.xr_session_settings
    landmark_active = VRLandmark.get_active_landmark(context)

    # Update the anchor object to the (new) camera of this landmark.
    session_settings.base_pose_object = landmark_active.base_pose_camera


def xr_landmark_active_base_pose_location_update(self, context):
    session_settings = context.window_manager.xr_session_settings
    landmark_active = VRLandmark.get_active_landmark(context)

    session_settings.base_pose_location = landmark_active.base_pose_location


def xr_landmark_active_base_pose_angle_update(self, context):
    session_settings = context.window_manager.xr_session_settings
    landmark_active = VRLandmark.get_active_landmark(context)

    session_settings.base_pose_angle = landmark_active.base_pose_angle


def xr_landmark_type_update(self, context):
    landmark_selected = VRLandmark.get_selected_landmark(context)
    landmark_active = VRLandmark.get_active_landmark(context)

    # Only update session settings data if the changed landmark is actually
    # the active one.
    if landmark_active == landmark_selected:
        xr_landmark_active_type_update(self, context)


def xr_landmark_camera_update(self, context):
    landmark_selected = VRLandmark.get_selected_landmark(context)
    landmark_active = VRLandmark.get_active_landmark(context)

    # Only update session settings data if the changed landmark is actually
    # the active one.
    if landmark_active == landmark_selected:
        xr_landmark_active_camera_update(self, context)


def xr_landmark_base_pose_location_update(self, context):
    landmark_selected = VRLandmark.get_selected_landmark(context)
    landmark_active = VRLandmark.get_active_landmark(context)

    # Only update session settings data if the changed landmark is actually
    # the active one.
    if landmark_active == landmark_selected:
        xr_landmark_active_base_pose_location_update(self, context)


def xr_landmark_base_pose_angle_update(self, context):
    landmark_selected = VRLandmark.get_selected_landmark(context)
    landmark_active = VRLandmark.get_active_landmark(context)

    # Only update session settings data if the changed landmark is actually
    # the active one.
    if landmark_active == landmark_selected:
        xr_landmark_active_base_pose_angle_update(self, context)


def xr_landmark_camera_object_poll(self, object):
    return object.type == 'CAMERA'


def xr_landmark_active_update(self, context):
    wm = context.window_manager

    xr_landmark_active_type_update(self, context)
    xr_landmark_active_camera_update(self, context)
    xr_landmark_active_base_pose_location_update(self, context)
    xr_landmark_active_base_pose_angle_update(self, context)

    if wm.xr_session_state:
        wm.xr_session_state.reset_to_base_pose(context)


class VRLandmark(bpy.types.PropertyGroup):
    name: bpy.props.StringProperty(
        name="VR Landmark",
        default="Landmark"
    )
    type: bpy.props.EnumProperty(
        name="Type",
        items=[
            ('SCENE_CAMERA', "Scene Camera",
             "Use scene's currently active camera to define the VR view base "
             "location and rotation"),
            ('USER_CAMERA', "Custom Camera",
             "Use an existing camera to define the VR view base location and "
             "rotation"),
            # Custom base poses work, but it's uncertain if they are really
            # needed. Disabled for now.
            # ('CUSTOM', "Custom Pose",
            #  "Allow a manually definied position and rotation to be used as "
            #  "the VR view base pose"),
        ],
        default='SCENE_CAMERA',
        update=xr_landmark_type_update,
    )
    base_pose_camera: bpy.props.PointerProperty(
        name="Camera",
        type=bpy.types.Object,
        poll=xr_landmark_camera_object_poll,
        update=xr_landmark_camera_update,
    )
    base_pose_location: bpy.props.FloatVectorProperty(
        name="Base Pose Location",
        subtype='TRANSLATION',
        update=xr_landmark_base_pose_location_update,
    )

    base_pose_angle: bpy.props.FloatProperty(
        name="Base Pose Angle",
        subtype='ANGLE',
        update=xr_landmark_base_pose_angle_update,
    )

    @staticmethod
    def get_selected_landmark(context):
        scene = context.scene
        landmarks = scene.vr_landmarks

        return (
            None if (len(landmarks) <
                     1) else landmarks[scene.vr_landmarks_selected]
        )

    @staticmethod
    def get_active_landmark(context):
        scene = context.scene
        landmarks = scene.vr_landmarks

        return (
            None if (len(landmarks) <
                     1) else landmarks[scene.vr_landmarks_active]
        )


class VIEW3D_UL_vr_landmarks(bpy.types.UIList):
    def draw_item(self, context, layout, _data, item, icon, _active_data,
                  _active_propname, index):
        landmark = item
        landmark_active_idx = context.scene.vr_landmarks_active

        layout.emboss = 'NONE'

        layout.prop(landmark, "name", text="")

        icon = (
            'RADIOBUT_ON' if (index == landmark_active_idx) else 'RADIOBUT_OFF'
        )
        props = layout.operator(
            "view3d.vr_landmark_activate", text="", icon=icon)
        props.index = index


class VIEW3D_PT_vr_landmarks(bpy.types.Panel):
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = "VR"
    bl_label = "Landmarks"
    bl_options = {'DEFAULT_CLOSED'}

    def draw(self, context):
        layout = self.layout
        scene = context.scene
        landmark_selected = VRLandmark.get_selected_landmark(context)

        layout.use_property_split = True
        layout.use_property_decorate = False  # No animation.

        row = layout.row()

        row.template_list("VIEW3D_UL_vr_landmarks", "", scene, "vr_landmarks",
                          scene, "vr_landmarks_selected", rows=3)

        col = row.column(align=True)
        col.operator("view3d.vr_landmark_add", icon='ADD', text="")
        col.operator("view3d.vr_landmark_remove", icon='REMOVE', text="")

        if landmark_selected:
            layout.prop(landmark_selected, "type")

            if landmark_selected.type == 'USER_CAMERA':
                layout.prop(landmark_selected, "base_pose_camera")
            # elif landmark_selected.type == 'CUSTOM':
            #     layout.prop(landmark_selected,
            #                 "base_pose_location", text="Location")
            #     layout.prop(landmark_selected,
            #                 "base_pose_angle", text="Angle")


class VIEW3D_PT_vr_session_view(bpy.types.Panel):
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = "VR"
    bl_label = "View"

    def draw(self, context):
        layout = self.layout
        session_settings = context.window_manager.xr_session_settings

        layout.use_property_split = True
        layout.use_property_decorate = False  # No animation.

        col = layout.column(align=True, heading="Show")
        col.prop(session_settings, "show_floor", text="Floor")
        col.prop(session_settings, "show_annotation", text="Annotations")

        col = layout.column(align=True)
        col.prop(session_settings, "clip_start", text="Clip Start")
        col.prop(session_settings, "clip_end", text="End")


class VIEW3D_PT_vr_session(bpy.types.Panel):
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = "VR"
    bl_label = "VR Session"

    def draw(self, context):
        layout = self.layout
        session_settings = context.window_manager.xr_session_settings

        layout.use_property_split = False

        is_session_running = bpy.types.XrSessionState.is_running(context)

        # Using SNAP_FACE because it looks like a stop icon -- I shouldn't
        # have commit rights...
        toggle_info = (
            ("Start VR Session", 'PLAY') if not is_session_running else (
                "Stop VR Session", 'SNAP_FACE')
        )
        layout.operator("wm.xr_session_toggle",
                        text=toggle_info[0], icon=toggle_info[1])

        layout.separator()

        layout.prop(session_settings, "use_positional_tracking")


class VIEW3D_OT_vr_landmark_add(bpy.types.Operator):
    bl_idname = "view3d.vr_landmark_add"
    bl_label = "Add VR Landmark"
    bl_description = "Add a new VR landmark to the list and select it"
    bl_options = {'UNDO', 'REGISTER'}

    def execute(self, context):
        scene = context.scene
        landmarks = scene.vr_landmarks

        landmarks.add()

        # select newly created set
        scene.vr_landmarks_selected = len(landmarks) - 1

        return {'FINISHED'}


class VIEW3D_OT_vr_landmark_remove(bpy.types.Operator):
    bl_idname = "view3d.vr_landmark_remove"
    bl_label = "Remove VR Landmark"
    bl_description = "Delete the selected VR landmark from the list"
    bl_options = {'UNDO', 'REGISTER'}

    def execute(self, context):
        scene = context.scene
        landmarks = scene.vr_landmarks

        if len(landmarks) > 1:
            landmark_selected_idx = scene.vr_landmarks_selected
            landmarks.remove(landmark_selected_idx)

            scene.vr_landmarks_selected -= 1

        return {'FINISHED'}


class VIEW3D_OT_vr_landmark_activate(bpy.types.Operator):
    bl_idname = "view3d.vr_landmark_activate"
    bl_label = "Activate VR Landmark"
    bl_description = "Change to the selected VR landmark from the list"
    bl_options = {'UNDO', 'REGISTER'}

    index: IntProperty(
        name="Index",
        options={'HIDDEN'},
    )

    def execute(self, context):
        scene = context.scene

        if self.index >= len(scene.vr_landmarks):
            return {'CANCELLED'}

        scene.vr_landmarks_active = (
            self.index if self.properties.is_property_set(
                "index") else scene.vr_landmarks_selected
        )

        return {'FINISHED'}


class VIEW3D_PT_vr_viewport_feedback(bpy.types.Panel):
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = "VR"
    bl_label = "Viewport Feedback"
    bl_options = {'DEFAULT_CLOSED'}

    def draw(self, context):
        layout = self.layout
        view3d = context.space_data

        col = layout.column(align=True)
        col.label(icon='ERROR', text="Note:")
        col.label(text="Settings here may have a significant")
        col.label(text="performance impact!")

        layout.separator()

        layout.prop(view3d.shading, "vr_show_virtual_camera")
        layout.prop(view3d, "mirror_xr_session")


class VIEW3D_GT_vr_camera_cone(Gizmo):
    bl_idname = "VIEW_3D_GT_vr_camera_cone"

    aspect = 1.0, 1.0

    def draw(self, context):
        import bgl

        if not hasattr(self, "frame_shape"):
            aspect = self.aspect

            frame_shape_verts = (
                (-aspect[0], -aspect[1], -1.0),
                (aspect[0], -aspect[1], -1.0),
                (aspect[0], aspect[1], -1.0),
                (-aspect[0], aspect[1], -1.0),
            )
            lines_shape_verts = (
                (0.0, 0.0, 0.0),
                frame_shape_verts[0],
                (0.0, 0.0, 0.0),
                frame_shape_verts[1],
                (0.0, 0.0, 0.0),
                frame_shape_verts[2],
                (0.0, 0.0, 0.0),
                frame_shape_verts[3],
            )

            self.frame_shape = self.new_custom_shape(
                'LINE_LOOP', frame_shape_verts)
            self.lines_shape = self.new_custom_shape(
                'LINES', lines_shape_verts)

        # Ensure correct GL state (otherwise other gizmos might mess that up)
        bgl.glLineWidth(1)
        bgl.glEnable(bgl.GL_BLEND)

        self.draw_custom_shape(self.frame_shape)
        self.draw_custom_shape(self.lines_shape)


class VIEW3D_GGT_vr_viewer_pose(GizmoGroup):
    bl_idname = "VIEW3D_GGT_vr_viewer_pose"
    bl_label = "VR Viewer Pose Indicator"
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'WINDOW'
    bl_options = {'3D', 'PERSISTENT', 'SCALE', 'VR_REDRAWS'}

    @classmethod
    def poll(cls, context):
        view3d = context.space_data
        return (
            view3d.shading.vr_show_virtual_camera and
            bpy.types.XrSessionState.is_running(context) and
            not view3d.mirror_xr_session
        )

    @staticmethod
    def _get_viewer_pose_matrix(context):
        from mathutils import Matrix, Quaternion

        wm = context.window_manager

        loc = wm.xr_session_state.viewer_pose_location
        rot = wm.xr_session_state.viewer_pose_rotation

        rotmat = Matrix.Identity(3)
        rotmat.rotate(rot)
        rotmat.resize_4x4()
        transmat = Matrix.Translation(loc)

        return transmat @ rotmat

    def setup(self, context):
        gizmo = self.gizmos.new(VIEW3D_GT_vr_camera_cone.bl_idname)
        gizmo.aspect = 1 / 3, 1 / 4

        gizmo.color = gizmo.color_highlight = 0.2, 0.6, 1.0
        gizmo.alpha = 1.0

        self.gizmo = gizmo

    def draw_prepare(self, context):
        self.gizmo.matrix_basis = self._get_viewer_pose_matrix(context)


classes = (
    VIEW3D_PT_vr_session,
    VIEW3D_PT_vr_session_view,
    VIEW3D_PT_vr_landmarks,
    VIEW3D_PT_vr_viewport_feedback,

    VRLandmark,
    VIEW3D_UL_vr_landmarks,

    VIEW3D_OT_vr_landmark_add,
    VIEW3D_OT_vr_landmark_remove,
    VIEW3D_OT_vr_landmark_activate,

    VIEW3D_GT_vr_camera_cone,
    VIEW3D_GGT_vr_viewer_pose,
)

# this class manages VR navigation using VR controllers.
# the Fly navigation allows user to manipulate the scene using the left controller (world scale is modified using left joystick).
# The Bi-manual navigation is the standard VR navigation with 2 controllers.
class Navigation:
    griped = False
    init_pivot_invert_matrix = None
    init_world_matrix = None
    scale = 1.0
    init_distance = 0.0
    navigation_mode = "Bimanual"
    fixedScaleFactor = 1.05

    hysteresis_min = 0.3
    hysteresis_max = 0.5

    def navigate(self):
        try:
            # trys to get an openxr session state
            wm = bpy.context.window_manager
            if not wm.xr_session_state:
                return

            session_state = wm.xr_session_state        
            if self.navigation_mode == "Fly":
                self.navigation_fly(session_state)

            if self.navigation_mode == "Bimanual":
                self.navigation_bimanual(session_state)
        except:
            # openxr is not ready
            pass

    def set_mode(self, mode):
        self.griped = False
        self.navigation_mode = mode

    def navigation_fly(self, session_state):
        grip = session_state.left_grip_value
        justGriped = False
        if not self.griped and grip > self.hysteresis_max:
            # left grip is pressed
            self.griped = True
            justGriped = True
        else:
            if self.griped and grip < self.hysteresis_min:
                # left grip is released
                self.griped = False

        if self.griped:
            # get left controller matrix
            controller_translate = session_state.left_controller_location
            controller_rotate = session_state.left_controller_rotation
            controller_matrix = Matrix.Translation(controller_translate) @ controller_rotate.to_matrix().to_4x4()

        if justGriped:
            # get init pivot and world matrix
            self.scale = 1.0
            self.init_pivot_invert_matrix = controller_matrix.inverted()
            self.init_world_matrix = Matrix.Translation(session_state.world_location) @ session_state.world_rotation.to_matrix().to_4x4() @ Matrix.Scale(session_state.world_scale, 4)

        if self.griped:
            # compute scale from joystick y value
            joyvalue = session_state.left_thumbstick_y
            if joyvalue > self.hysteresis_max:
                self.scale *= self.fixedScaleFactor
            if joyvalue < -self.hysteresis_min:
                self.scale /= self.fixedScaleFactor

            # compute new fake world matrix
            scale_matrix = Matrix.Scale(self.scale, 4)
            transformedWorld = controller_matrix @ scale_matrix @ self.init_pivot_invert_matrix @ self.init_world_matrix

            # apply to fake world transform
            session_state.world_location = transformedWorld.to_translation()
            session_state.world_rotation = transformedWorld.to_quaternion()
            session_state.world_scale = transformedWorld.to_scale().x

    def navigation_bimanual(self, session_state):
        # get controllers grip values
        left_grip = session_state.left_grip_value
        right_grip = session_state.right_grip_value
        justGriped = False
        if not self.griped and left_grip > self.hysteresis_max and right_grip > self.hysteresis_max:
            # if both grips are pressed
            self.griped = True
            justGriped = True
        else:
            if self.griped and (left_grip < self.hysteresis_min or right_grip < self.hysteresis_min):
                # if both grips are released
                self.griped = False

        if self.griped:
            # get position of the center of the controllers
            center_position = (session_state.left_controller_location + session_state.right_controller_location) * 0.5
            # get vector from left controller to right controller
            delta = session_state.right_controller_location - session_state.left_controller_location
            # compute a controller matrix centerd between the controllers 
            # and oriented by projecting vector from left controller to right controller to a horizontal plane 
            right = delta.normalized()
            up = Vector((0,0,1))
            forward = Vector.cross(up, right)
            right = Vector.cross(forward, up)
            # create pivot matrix from this base
            pivot_matrix = Matrix()
            pivot_matrix.col[0] = Vector((right.x, right.y, right.z, 0.0))
            pivot_matrix.col[1] = Vector((forward.x, forward.y, forward.z, 0.0))
            pivot_matrix.col[2] = Vector((up.x, up.y, up.z, 0.0))
            pivot_matrix.col[3] = Vector((center_position.x, center_position.y, center_position.z, 1.0))

        if justGriped:
            # get init pivot and world matrix
            self.griped = True
            self.scale = 1.0
            self.init_distance = delta.length
            self.init_pivot_invert_matrix = pivot_matrix.inverted()
            self.init_world_matrix = Matrix.Translation(session_state.world_location) @ session_state.world_rotation.to_matrix().to_4x4() @ Matrix.Scale(session_state.world_scale, 4)

        if self.griped:
            # compute scale
            distance = delta.length
            self.scale = distance / self.init_distance

            # compute new fake world matrix
            scale_matrix = Matrix.Scale(self.scale, 4)
            transformedWorld = pivot_matrix @ scale_matrix @ self.init_pivot_invert_matrix @ self.init_world_matrix

            # apply to fake world transform
            session_state.world_location = transformedWorld.to_translation()
            session_state.world_rotation = transformedWorld.to_quaternion()
            session_state.world_scale = transformedWorld.to_scale().x

navigation = Navigation()

@persistent
def controller_handler():
    global navigation
    navigation.navigate()
    return 0.0

@persistent
def register_controller_handler(context: bpy.context):
    bpy.app.timers.register(controller_handler)    

def register():
    if not bpy.app.build_options.xr_openxr:
        return

    for cls in classes:
        bpy.utils.register_class(cls)

    bpy.types.Scene.vr_landmarks = CollectionProperty(
        name="Landmark",
        type=VRLandmark,
    )
    bpy.types.Scene.vr_landmarks_selected = IntProperty(
        name="Selected Landmark"
    )
    bpy.types.Scene.vr_landmarks_active = IntProperty(
        update=xr_landmark_active_update,
    )
    # View3DShading is the only per 3D-View struct with custom property
    # support, so "abusing" that to get a per 3D-View option.
    bpy.types.View3DShading.vr_show_virtual_camera = BoolProperty(
        name="Show VR Camera"
    )

    bpy.app.handlers.load_post.append(ensure_default_vr_landmark)
    bpy.app.handlers.load_post.append(register_controller_handler)
    bpy.app.timers.register(controller_handler)

def unregister():
    if not bpy.app.build_options.xr_openxr:
        return

    for cls in classes:
        bpy.utils.unregister_class(cls)

    del bpy.types.Scene.vr_landmarks
    del bpy.types.Scene.vr_landmarks_selected
    del bpy.types.Scene.vr_landmarks_active
    del bpy.types.View3DShading.vr_show_virtual_camera

    bpy.app.timers.unregister(controller_handler)
    bpy.app.handlers.load_post.remove(register_controller_handler)
    bpy.app.handlers.load_post.remove(ensure_default_vr_landmark)

if __name__ == "__main__":
    register()