tutorials/02_randomization.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Tutorial02 - Randomization\n",
                "\n",
                "## Overview\n",
                "\n",
                "🚀 By the last tutorial, you have learned the concept of ``Actor``, ``Level`` and ``Sequence``. In this tutorial, you will learn how to add random transform keys to an ``Actor`` in order to render multiple poses of it. By the end of this tutorial, you will be able to:\n",
                "\n",
                "- Initialize XRFeitoria\n",
                "- Import ``Actor`` and label it by ``stencil value``\n",
                "- Set the scale of ``Actor`` in the ``Level``\n",
                "- Create a `Sequence` for rendering and adding transform keys to `Actor`\n",
                "- Add a camera in the `Sequence`\n",
                "- Render images and annotations"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 1. Initialization\n",
                "\n",
                "Install the following packages that will be used in this tutorial:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "%pip install objaverse\n",
                "%pip install scipy"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Then, similar to [Tutorial01](./01_get_started.ipynb), specify your engine path and initialize XRFeitoria."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import xrfeitoria as xf"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Replace with your executable path\n",
                "engine_exec_path = 'C:/Program Files/Blender Foundation/Blender 3.3/blender.exe'\n",
                "# engine_exec_path = 'C:/Program Files/Epic Games/UE_5.2/Engine/Binaries/Win64/UnrealEditor-Cmd.exe'"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from pathlib import Path\n",
                "\n",
                "exec_path_stem = Path(engine_exec_path).stem.lower()\n",
                "if 'blender' in exec_path_stem:\n",
                "    # Open Blender\n",
                "    render_engine = 'blender'\n",
                "    xf_runner = xf.init_blender(exec_path=engine_exec_path, background=False, new_process=True)\n",
                "elif 'unreal' in exec_path_stem:\n",
                "    # Unreal Engine requires a project to be opened\n",
                "    # Here we use a sample project, which is downloaded from the following link\n",
                "    # You can also use your own project\n",
                "    import shutil\n",
                "\n",
                "    from xrfeitoria.utils.downloader import download\n",
                "    unreal_project_zip = download(url='https://openxrlab-share.oss-cn-hongkong.aliyuncs.com/xrfeitoria/tutorials/unreal_project/UE_Sample.zip', \n",
                "                                    dst_dir=\"./tutorial02/assets/\")\n",
                "    shutil.unpack_archive(filename=unreal_project_zip, extract_dir='./tutorial02/assets/')\n",
                "\n",
                "    # Open Unreal Engine\n",
                "    render_engine = 'unreal'\n",
                "    xf_runner = xf.init_unreal(exec_path=engine_exec_path, \n",
                "                                background=False, \n",
                "                                new_process=True, \n",
                "                                project_path='./tutorial02/assets/UE_sample/UE_sample.uproject')"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "✨ Now you can see a new Blender/Unreal Engine process has started."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 2. Import meshes\n",
                "\n",
                "Download some meshes from [Objaverse](https://objaverse.allenai.org/objaverse-1.0/)."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import objaverse\n",
                "\n",
                "objects = objaverse.load_objects(\n",
                "    uids=['eb0807309530496aaab9dcff67bf5c31',\n",
                "          'b4065dd5ce9d46be90db3e1f3e4b9cc1',\n",
                "          '0176be079c2449e7aaebfb652910a854',\n",
                "          'f130ebeb60f24ed8bd3714a7ed3ba280',\n",
                "          '289a2221178843a78ad433705555e16a',\n",
                "          'b7f7ab9bf7244c3a8851bae3fb0bf741',\n",
                "        ],\n",
                "    download_processes=1\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Import the meshes to create ``Actor`` instances in the ``Level``. \n",
                "\n",
                "Here we set different ``stencil_value`` for each ``Actor``. The [stencil value](https://xrfeitoria.readthedocs.io/en/latest/faq.html#what-is-stencil-value) is used to distinguish different Actors when rendering segmentation masks."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "actors = []\n",
                "for idx, file_path in enumerate(objects.values()):\n",
                "    actor = xf_runner.Actor.import_from_file(\n",
                "        file_path=file_path, \n",
                "        stencil_value=(idx+1)*10\n",
                "    )\n",
                "    actors.append(actor)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Switch to the engine window, and you can see the meshes has been imported.\n",
                "\n",
                "Then, we adjust the scale of the Actors to make their sizes equal to 0.3m."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "actor_size = 0.3\n",
                "for actor in actors:\n",
                "    _scale = actor_size / max(actor.dimensions)\n",
                "    actor.scale = (_scale, _scale, _scale)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Now they look like:\n",
                "\n",
                "- Blender\n",
                "![](https://openxrlab-share.oss-cn-hongkong.aliyuncs.com/xrfeitoria/tutorials/imgs/02/blender/scale_the_actors.png)\n",
                "\n",
                "- Unreal Engine\n",
                "![](https://openxrlab-share.oss-cn-hongkong.aliyuncs.com/xrfeitoria/tutorials/imgs/02/unreal/scale_the_actors.png)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "If you use ``Unreal Engine``, the ``Level`` should be saved after been modified."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# save the level\n",
                "if render_engine == 'unreal':\n",
                "    xf_runner.utils.save_current_level()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 3. Add a sequence for rendering and adding transform keys\n",
                "\n",
                "``Sequence`` is a multifunctional class in XRFeitoria. It can be used for:\n",
                "- rendering\n",
                "- adding transform keys\n",
                "- grouping different objects.\n",
                "\n",
                "Here, we use it for ``rendering`` and ``adding transform keys``."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "> ``Transform keys`` record the transformation (location, rotation and scale) of an ``Actor`` or a ``Camera`` at specific frames, and the transformation between two adjacent keys will be interpolated by the specified interpolation method. By adding transform keys, you can render multiple poses of an ``Actor``.\n",
                "> \n",
                "> In ``XRFeitoria``, transform keys are always stored in a list, and the members of the list are [SequenceTransformKey](https://xrfeitoria.readthedocs.io/en/latest/apis/generated/xrfeitoria.data_structure.models.SequenceTransformKey.html) object.\n",
                "\n",
                "Firstly, we randomly generate some transform keys for each actors."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import random\n",
                "\n",
                "from loguru import logger\n",
                "from scipy.stats import qmc\n",
                "\n",
                "from xrfeitoria.data_structure.models import SequenceTransformKey as SeqTransKey\n",
                "\n",
                "# Set the number of frames\n",
                "frame_num = 10\n",
                "\n",
                "# Use a dictionary to store the transform keys of all actors\n",
                "all_actors_transform_keys = {actor: [] for actor in actors}\n",
                "\n",
                "# Iterate over all frames\n",
                "for i in range(frame_num):\n",
                "\n",
                "    # Generate random locations by Poisson Disk Sampling\n",
                "    # The minimum distance between two actors is set to be `actor_size` we defined before\n",
                "    posson_engine = qmc.PoissonDisk(d=2, radius=actor_size)\n",
                "    sample_location = posson_engine.random(len(actors))\n",
                "    \n",
                "    # Set the transform keys for each actor\n",
                "    for actor_idx, actor in enumerate(actors):\n",
                "        actor_scale = actor.scale\n",
                "        \n",
                "        # Get the location from the samples generated by Poisson Disk Sampling\n",
                "        random_location = (sample_location[actor_idx][0], \n",
                "                           0.0, \n",
                "                           sample_location[actor_idx][1])\n",
                "        \n",
                "        # Generate random rotations\n",
                "        random_rotation = (random.random() * 360.0,\n",
                "                           random.random() * 360.0,\n",
                "                           random.random() * 360.0)\n",
                "        \n",
                "        # Generate random scales\n",
                "        scale = random.uniform(0.5, 1.0)\n",
                "        random_scale = (scale * actor_scale[0],\n",
                "                        scale * actor_scale[1],\n",
                "                        scale * actor_scale[2])\n",
                "        \n",
                "        # Save the transform keys\n",
                "        all_actors_transform_keys[actor].append(\n",
                "            SeqTransKey(\n",
                "                frame=i,\n",
                "                location=random_location,\n",
                "                rotation=random_rotation,\n",
                "                scale=random_scale,\n",
                "                interpolation='AUTO',\n",
                "            )            \n",
                "        )\n",
                "    logger.info(f'Generated transform keys of frame {i}.')"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Then we create a ``Sequence`` to apply the transform keys to the actors and render the images."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from xrfeitoria.data_structure.models import RenderPass\n",
                "\n",
                "# Use the `with` statement to create a sequence, and it will be automatically close the sequence after the code block is executed.\n",
                "# The argument `seq_length` controls the number of frames to be rendered. \n",
                "sequence_name = 'MySequence'\n",
                "with xf_runner.Sequence.new(seq_name=sequence_name, seq_length=frame_num, replace=True) as seq:\n",
                "    ##############################\n",
                "    ##### Add transform keys #####\n",
                "    ##############################\n",
                "\n",
                "    # The function `use_actor_with_keys` sets transform keys for the actor in the sequence.\n",
                "    # The transform keys are only stored in the sequence.\n",
                "    # When the sequence is closed, the actor will be restored to its original state(without transform keys).\n",
                "    for actor, keys in all_actors_transform_keys.items():\n",
                "        seq.use_actor_with_keys(actor=actor, transform_keys=keys)\n",
                "\n",
                "    #####################\n",
                "    ##### Rendering #####\n",
                "    #####################\n",
                "\n",
                "    # Add a camera and make it look at the specified location\n",
                "    camera_location = (0.5, -3.0, 0.5)\n",
                "    camera_rotation = xf_runner.utils.get_rotation_to_look_at(location=camera_location, target=(0.5, 0.0, 0.5))\n",
                "    camera = seq.spawn_camera(location=camera_location, rotation=camera_rotation, fov=45)\n",
                "\n",
                "    # Add a render job to renderer\n",
                "    # In render job, you can specify the output path, resolution, render passes, etc.\n",
                "    # The ``output_path`` is the path to save the rendered data.\n",
                "    # The ``resolution`` is the resolution of the rendered image.\n",
                "    # The ``render_passes`` define what kind of data you want to render, such as img, depth, normal, etc.\n",
                "    # and what kind of format you want to save, such as png, jpg, exr, etc.\n",
                "    seq.add_to_renderer(\n",
                "        output_path=f'./tutorial02/outputs/{render_engine}/',\n",
                "        resolution=(1280, 720),\n",
                "        render_passes=[RenderPass('img', 'png'),\n",
                "                       RenderPass('mask', 'exr'),\n",
                "                       RenderPass('normal', 'exr'),\n",
                "                       RenderPass('diffuse', 'exr')],\n",
                "    )"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 4. Render\n",
                "\n",
                "The following code renders all the render jobs and save the images to the ``output_path`` set in ``seq.add_to_renderer`` above."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Render\n",
                "xf_runner.render()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Check the ``output_path``, and you can see the rendered images and their annotations. Visualize the images and annotations by the following code."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import matplotlib.pyplot as plt\n",
                "\n",
                "from xrfeitoria.utils.viewer import Viewer\n",
                "\n",
                "xf_viewer = Viewer(sequence_dir=f'./tutorial02/outputs/{render_engine}/{sequence_name}/')\n",
                "\n",
                "for i in range(frame_num):\n",
                "    img = xf_viewer.get_img(camera_name=camera.name, frame=i)\n",
                "    mask = xf_viewer.get_mask(camera_name=camera.name, frame=i)\n",
                "    normal = xf_viewer.get_normal(camera_name=camera.name, frame=i)\n",
                "    diffuse = xf_viewer.get_diffuse(camera_name=camera.name, frame=i)\n",
                "\n",
                "    plt.figure(figsize=(20, 20))\n",
                "\n",
                "    plt.subplot(1, 4, 1)\n",
                "    plt.imshow(img)\n",
                "    plt.axis('off')\n",
                "    plt.title('img')\n",
                "\n",
                "    plt.subplot(1, 4, 2)\n",
                "    plt.imshow(mask)\n",
                "    plt.axis('off')\n",
                "    plt.title('mask')\n",
                "\n",
                "    plt.subplot(1, 4, 3)\n",
                "    plt.imshow(normal)\n",
                "    plt.axis('off')\n",
                "    plt.title('normal')\n",
                "\n",
                "    plt.subplot(1, 4, 4)\n",
                "    plt.imshow(diffuse)\n",
                "    plt.axis('off')\n",
                "    plt.title('diffuse')"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "> Hint: When using Unreal Engine, if the images of the mask look weird, try running the notebook again."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Finally, close the engine by:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Close the engine\n",
                "xf_runner.close()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Ref to [api docs](https://xrfeitoria.readthedocs.io/en/latest/apis/xrfeitoria.html), you can always use ``with`` statement to ensure the engine is closed when the codes are finished."
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "blender-pipeline",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.10.12"
        },
        "orig_nbformat": 4
    },
    "nbformat": 4,
    "nbformat_minor": 2
}