Capturing screenshots

While there are a number of additional features we could conceive adding to our Service, there is little more we need to add to have a functional viewport capture capability.

Despite having only provided placeholder functionality in the core of our service, we’ll be able to have image capture capability with little additional changes to the existing code we have implemented so far. This is a valuable advantage, as the care we have taken to follow best practices so far would otherwise have been of little return on investment.

Building on existing features

As you may have already noticed if you had the opportunity to use Omniverse applications before, there is already an option to capture screenshots present in the Edit > Capture Screenshot menu (and also available using the F10 key of the keyboard). In order to reuse this feature which already brings us very close to our intended goal, we will be wrapping this convenient feature under a utility function, which we’ll then only use to call in the function handler for our Service.

This will be a flexible solution for our Service moving forward too, as we can conceive of a number of additional desirable features we would wish to support in the future too, which we could develop in parallel to the version we could already ship to a number of Users for early feedback, before swapping to a newer version.

We could even conceive exposing different image capture implementations in the same Service, with our function handler accepting an optional argument that would select which of our capture features to execute. This would allow us to ship this candidate for an improved version of the Service to a number of Users, and use a feature flag to request their feedback on alternative methods so we could collect their insights about potential stability or performance issues before rolling out the feature to a larger audience in a future version of the Service.

First things first, let’s start by introducing a few utility methods to capture and serve images from a stored location. We will implement them independently from our Service stack, and take care of decoupling them from the message-handling section of our API in order to make it easier to evolve these sections independently of each other. This will not only make testing them easier, this will also promote code reuse across projects. Multiple benefits!

exts/omni.services.example.viewport.capture/omni/services/example/viewport/capture/utils.py

import asyncio
import os
import shutil
from typing import Optional, Tuple

import carb.settings
import carb.tokens

import omni.kit.actions.core
import omni.kit.app
import omni.usd


# Let's include a small utility method to facilitate obtaining the name of the extension our code is bundled with.
# While we could certainly store and share the `ext_id` provided to the `on_startup()` method of our Extension, this
# alternative method of obtaining the name of our extension can also make our code more portable across projects, as it
# may allow you to keep your code changes located closer together and not have to spread them up to the main entrypoint
# of your extension.
def get_extension_name() -> str:
    """
    Return the name of the Extension where the module is defined.

    Args:
        None

    Returns:
        str: The name of the Extension where the module is defined.

    """
    extension_manager = omni.kit.app.get_app().get_extension_manager()
    extension_id = extension_manager.get_extension_id_by_module(__name__)
    extension_name = extension_id.split("-")[0]
    return extension_name


# Building on the utility method just above, this helper method helps us retrieve the path where captured images are
# served from the web server, so they can be presented to clients over the network.
def get_captured_image_path() -> str:
    """
    Return the path where the captured images can be retrieved from the server, in the `/{url_prefix}/{capture_path}`
    format.

    Args:
        None

    Returns:
        str: The path where the captured images can be retrieved from the server.

    """
    extension_name = get_extension_name()
    settings = carb.settings.get_settings()
    url_prefix = settings.get_as_string(f"exts/{extension_name}/url_prefix")
    capture_path = settings.get_as_string(f"exts/{extension_name}/capture_path")

    captured_images_path = f"{url_prefix}{capture_path}"
    return captured_images_path


# In a similar fashion to the utility method above, this helper method helps us retrieve the path on disk where the
# captured images are stored on the server. This makes it possible to map this storage location known to the server to a
# publicly-accessible location on the server, from which clients will be able to fetch the captured images once their
# web-friendly names have been communicated to clients through our Service's response.
def get_captured_image_directory() -> str:
    """
    Return the location on disk where the captured images will be stored, and from which they will be served by the web
    server after being mounted. In order to avoid growing the size of this static folder indefinitely, images will be
    stored under the `${temp}` folder of the Omniverse application folder, which is emptied when the application is shut
    down.

    Args:
        None

    Returns:
        str: The location on disk where the captured images will be stored.

    """
    extension_name = _get_extension_name()
    capture_directory_name = carb.settings.get_settings().get_as_string(f"exts/{extension_name}/capture_directory")
    temp_kit_directory = carb.tokens.get_tokens_interface().resolve("${temp}")
    captured_stage_images_directory = os.path.join(temp_kit_directory, capture_directory_name)
    return captured_stage_images_directory


# This is the main utility method of our collection so far. This small helper builds on the existing capability of the
# "Edit > Capture Screenshot" feature already available in the menu to capture an image from the Omniverse application
# currently running. Upon completion, the captured image is moved to the storage location that is mapped to a
# web-accessible path so that clients are able to retrieve the screenshot once they are informed of the image's unique
# name when our Service issues its response.
async def capture_viewport(usd_stage_path: str) -> Tuple[bool, Optional[str], Optional[str]]:
    """
    Capture the viewport, by executing the action already registered in the "Edit > Capture Screenshot" menu.

    Args:
        usd_stage_path (str): Path of the USD stage to open in the application's viewport.

    Returns:
        Tuple[bool, Optional[str], Optional[str]]: A tuple containing a flag indicating the success of the operation,
            the path of the captured image on the web server, along with an optional error message in case of error.

    """
    success: bool = omni.usd.get_context().open_stage(usd_stage_path)
    captured_image_path: Optional[str] = None
    error_message: Optional[str] = None

    if success:
        event = asyncio.Event()

        menu_action_success: bool = False
        capture_screenshot_filepath: Optional[str] = None
        def callback(success: bool, captured_image_path: str) -> None:
            nonlocal menu_action_success, capture_screenshot_filepath
            menu_action_success = success
            capture_screenshot_filepath = captured_image_path

            event.set()

        omni.kit.actions.core.execute_action("omni.kit.menu.edit", "capture_screenshot", callback)
        await event.wait()
        await asyncio.sleep(delay=1.0)

        if menu_action_success:
            # Move the screenshot to the location from where it can be served over the network:
            destination_filename = os.path.basename(capture_screenshot_filepath)
            destination_filepath = os.path.join(get_captured_image_directory(), destination_filename)
            shutil.move(src=capture_screenshot_filepath, dst=destination_filepath)

            # Record the final location of the captured image, along with the status of the operation:
            captured_image_path = os.path.join(get_captured_image_path(), destination_filename)
            success = menu_action_success
    else:
        error_message = f"Unable to open stage \"{usd_stage_path}\"."

    return (success, captured_image_path, error_message)

Now that these utility methods are compartmentalized in their own independent module, away from the function-handling features of our Service, all we need to do to connect this live implementation of our feature is to call it from our Service:

exts/omni.services.example.viewport.capture/omni/services/example/viewport/capture/utils.py

# [...]

async def capture(request: ViewportCaptureRequestModel,) -> ViewportCaptureResponseModel:
-    # For now, let's just print incoming request to the log to confirm all components of our extension are properly
-    # wired together:
-    carb.log_warn(f"Received a request to capture an image of \"{request.usd_stage_path}\".")
+    success, captured_image_path, error_message = await capture_viewport(usd_stage_path=request.usd_stage_path)

    return ViewportCaptureResponseModel(
-       success=False,
-       captured_image_path=None,
-       error_message="Image not yet captured.",
+       success=success,
+       captured_image_path=captured_image_path,
+       error_message=error_message,
    )

Hopefully these few lines where we performed a change illustrate the benefits of the planning we performed initially, and the separation of concerns between the various sections of our implementations:

• The API section of our Service takes care of handling the requests from clients, routes messages to the appropriate delegates, and issues responses back to clients after performing work.

• The system can otherwise remain unaware of the way client requests or responses are formatted, allowing for flexibility in the way operations are performed. Changing the inner implementation of our operation will have limited impact on clients who may be using our Service. Similarly, Users of our Service can grow confident in the stability of our system, as changes we may be interested in performing over time will have limited impact on the way Users interface with our Service thanks to the fact that it does not leak its implementation details through the API.

Serving content to Users

One last operation we will perform in order to ensure the images we captured can be properly served back to clients through the network is to map a web-friendly storage folder on the server. Let’s update our extension entrypoint to register this mounted storage location, reusing the utility methods we conveniently shared earlier.

We’ve highlighted the changes we performed in this step to make it easier to reference the few features we added:

exts/omni.services.example.viewport.capture/omni/services/example/viewport/capture/extension.py

import os

from fastapi.staticfiles import StaticFiles

import carb

import omni.ext
from omni.services.core import main

# As most of the features of our API are implemented by the means of function handlers in the `/services` sub-folder,
# the main duty of our extension entrypoint is to register our Service's `@router` and dictate when its capability
# should be enabled or disabled under the guidance of the User or the dependency system.
from .services.capture import router

# For convenience, let's also reuse the utility methods we already created to handle and format the storage location of
# the captured images so they can be accessed by clients using the server, once API responses are issued from our
# Service:
from .utils import get_captured_image_directory, get_captured_image_path


# Any class derived from `omni.ext.IExt` in the top level module (defined in the `python.module` section of the
# `extension.toml` file) will be instantiated when the extension is enabled, and its `on_startup(ext_id)` method
# will be called. When disabled or when the application is shut down, its `on_shutdown()` will be called.
class ViewportCaptureExtension(omni.ext.IExt):
    """Sample extension illustrating registration of a service."""

    # `ext_id` is the unique identifier of the extension, containing its name and semantic version number. This
    # identifier can be used in conjunction with the Extension Manager to query for additional information, such
    # as the extension's location on the filesystem.
    def on_startup(self, ext_id: str) -> None:
        ext_name = ext_id.split("-")[0]

        # At this point, we register our Service's `router` under the prefix we gave our API using the settings system,
        # to facilitate its configuration and to ensure it is unique from all other extensions we may have enabled:
        url_prefix = carb.settings.get_settings().get_as_string(f"exts/{ext_name}/url_prefix")
        main.register_router(router=router, prefix=url_prefix, tags=["Viewport capture"],)

        # Proceed to create a temporary directory in the Omniverse application file hierarchy where captured stage
        # images will be stored, until the application is shut down:
        captured_stage_images_directory = get_captured_image_directory()
        if not os.path.exists(captured_stage_images_directory):
            os.makedirs(captured_stage_images_directory)

        # Register this location as a mount, so its content is served by the web server bundled with the Omniverse
        # application instance, thus making the captured image available on the network:
        main.register_mount(
            path=get_captured_image_path(),
            app=StaticFiles(directory=captured_stage_images_directory, html=True),
            name="captured-stage-images",
        )

    def on_shutdown(self) -> None:
        # When disabling the extension or shutting down the instance of the Omniverse application, let's make sure we
        # also deregister our Service's `router` in order to avoid our API being erroneously advertised as present as
        # part of the OpenAPI specification despite our handler function no longer being available:
        main.deregister_router(router=router)
        main.deregister_mount(path=get_captured_image_path())

Note

For additional information regarding how to execute this Service via code, consult the online documentation about the using the asynchronous request client.

Use cases for such use cases include:

• Executing this Omniverse Service from another Service, either on the same machine or from a remote one.

• Integrating the capabilities exposed by the Omniverse Service into another application.

• Executing this Omniverse Service from another technology stack, for which there may not be officially-supported bindings other than HTTP or other Omniverse Service Transports.

• etc.

« Previous section: Exploring the API	Next section: Providing unit tests »