scripts/ui/export.py (435 lines of code) (raw):
#!/usr/bin/env python3
# Copyright 2004-present Facebook. All Rights Reserved.
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
"""Export tab for the UI.
Defines the interaction model for the background tab in the UI. This tab cannot
be run in isolation and expects a very particular structure defined in dep.ui.
UI extensions to the tab can be made by modifying the QT structure and adding
corresponding functionality here.
Example:
To see an instance of Export, refer to the example in dep.py:
>>> export = Export(self) # self refers to the overall QT UI object
"""
import json
import os
import sys
from collections import OrderedDict
import pyvidia
dir_scripts = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
dir_root = os.path.dirname(dir_scripts)
sys.path.append(dir_root)
sys.path.append(os.path.join(dir_scripts, "render"))
import common
import dep_util
import scripts.render.config as config
import verify_data
from scripts.render.network import download
from scripts.util.system_util import (
get_flags_from_flagfile,
image_type_paths,
merge_lists,
)
class Export:
"""Tab in UI responsible for producing meshed binaries and other final output formats
(e.g. color and disparity equirects) and viewing them in corresponding PC and
headset displays.
Attributes:
dlg (App(QDialog)): Main dialog box for the app initialized by QT.
formats (dict[str, str]): Formats that can be rendered given the render configuration.
is_farm (bool): Whether or not this is a farm render.
is_host_linux_gpu (bool): Whether or not this is a Linux GPU machine or not.
is_local_non_win (bool): Whether this is a Windows machine or not.
is_viewer_rift_hidden (bool): Whether or not to display the "RiftViewer" button.
is_viewer_smr_hidden (bool): Whether or not to display the "SimpleMeshRenderer" button.
log_reader (LogReader(QtCore.QObject)): Object to interact with UI logging system.
parent (App(QDialog)): Object corresponding to the parent UI element.
path_rig_json (str): Path to the rig JSON.
tag (str): Semantic name for the tab.
"""
def __init__(self, parent):
"""Initializes the Export tab.
Args:
parent (App(QDialog)): Object corresponding to the parent UI element.
"""
self.parent = parent
self.tag = "export"
self.dlg = parent.dlg
common.init(self)
self.initialize_viewer_buttons()
def setup_aws_config(self):
"""Configures the AWS credentials."""
common.setup_aws_config(self)
def add_data_type_validators(self):
"""Adds validators for UI elements."""
dlg = self.dlg
elems = [
dlg.val_export_options_res,
dlg.dd_export_data_first,
dlg.dd_export_data_last,
]
for elem in elems:
dep_util.set_integer_validator(elem)
def initialize_viewer_buttons(self):
"""Sets up buttons to IPC callbacks on the host."""
btn_smr_onscreen = self.dlg.btn_export_data_smr_view
btn_riftviewer = self.dlg.btn_export_data_rift_view
host_os = self.parent.ui_flags.host_os
self.is_viewer_rift_hidden = False
self.is_viewer_smr_hidden = False
# Viewers only available if we have a local_bin path
if self.parent.ui_flags.local_bin == "":
btn_smr_onscreen.setEnabled(False)
btn_riftviewer.setEnabled(False)
self.is_viewer_rift_hidden = True
self.is_viewer_smr_hidden = True
# RiftViewer only available in a Windows host
if host_os != "OSType.WINDOWS":
btn_riftviewer.setEnabled(False)
self.is_viewer_rift_hidden = True
# SimpleMeshRenderer onscreen viewer only available in a non-Windows host
else:
btn_smr_onscreen.setEnabled(False)
self.is_viewer_smr_hidden = True
def setup_farm(self):
"""Sets up a Kubernetes farm for AWS renders."""
common.setup_farm(self)
def initialize_paths(self):
"""Initializes paths for scripts and flags."""
common.initialize_paths(self)
self.formats = {
"6DoF (Meshing)": "6dof_meshing",
"6DoF (Striping)": "6dof_striping",
}
# SimpleMeshRenderer only available to render when:
# - Cloud rendering (AWS)
# - Linux host with NVIDIA GPU
# - non-Windows host + given local_bin flag
host_os = self.parent.ui_flags.host_os
is_linux = host_os == "OSType.LINUX"
is_non_windows = host_os != "OSType.WINDOWS"
has_local_bin = self.parent.ui_flags.local_bin != ""
self.is_host_linux_gpu = is_linux and (pyvidia.get_nvidia_device() is not None)
self.is_local_non_win = is_non_windows and has_local_bin
if self.parent.is_aws or self.is_host_linux_gpu or self.is_local_non_win:
self.formats.update(
{
"Equirect color": "eqrcolor",
"Equirect disparity": "eqrdisp",
"Cubemap color": "cubecolor",
"Cubemap disparity": "cubedisp",
"180 stereo left-right": "lr180",
"360 stereo top-bottom": "tbstereo",
"3DoF top-bottom": "tb3dof",
}
)
def switch_ui_elements_for_processing(self, state):
"""Switches element interaction when processing.
Args:
state (bool): State to which elements should be changed to (i.e.
True for enabled, False for disabled)
"""
gb = self.dlg.gb_export_data
common.switch_ui_elements_for_processing(self, gb, state)
def sync_with_s3(self):
"""Syncs data available locally with the S3 bucket."""
if not self.is_farm: # rendered locally
gb = self.dlg.gb_export_data
subdirs = [
self.parent.path_flags,
image_type_paths["video_bin"],
image_type_paths["video_fused"],
]
common.sync_with_s3(self, gb, subdirs)
def on_process_finished(self, exitCode, exitStatus, p_id):
"""Callback event handler for a process completing.
Args:
exitCode (int): Return code from running the process.
exitStatus (str): Description message of process exit.
p_id (str): PID of completed process.
"""
common.on_process_finished(self, p_id)
def setup_logs(self):
"""Sets up logging system for dialog on the current tab."""
self.log_reader = common.setup_logs(self)
def set_default_top_level_paths(self, mkdirs=False):
"""Defines class referenceable attributes for paths. See common for full set
of definitions.
Args:
mkdirs (bool, optional): Whether or not to make the defined directories.
"""
verify_data.set_default_top_level_paths(self, mkdirs)
def get_ec2_instance_types(self):
"""Gets valid instances that can be spawned for a job.
Returns:
list[str]: Names of valid instance types.
"""
if self.get_format().startswith("6dof"):
return self.ec2_instance_types_cpu
else:
return self.ec2_instance_types_gpu
def get_valid_types(self):
"""Checks which render types are valid.
Returns:
list[str]: Subset of ("background", "video") depending on if the corresponding
inputs are present.
"""
# We need both full-size color and disparity for a type to be valid
type_paths = OrderedDict(
{
"color": [self.path_video_color, self.path_video_disparity],
"background_color": [self.path_bg_color, self.path_bg_disparity],
}
)
ps = []
for type, paths in type_paths.items():
if all(dep_util.check_image_existence(p) != "" for p in paths):
ps.append(type)
return ps
def get_frame_names(self):
"""Finds all the frames in a local directory.
Returns:
list[str]: Sorted list of frame names in the directory.
"""
s = "_s3" if self.is_aws and self.is_farm else ""
frames_bin = getattr(self.parent, f"frames_bin{s}", None)
data_type = self.dlg.dd_export_data_type.currentText()
if not data_type:
return frames_bin
t = "bg" if data_type == "background_color" else "video"
frames_color = getattr(self.parent, f"frames_{t}_color{s}", None)
frames_disp = getattr(self.parent, f"frames_{t}_disparity{s}", None)
return sorted(merge_lists(frames_color, frames_disp, frames_bin))
def get_files(self, tag):
"""Retrieves file names corresponding to the desired attribute.
Args:
tag (str): Semantic name for the attribute.
Returns:
list[str]: List of file names.
Raises:
Exception: If a tag is requested with no associated files.
"""
if tag in ["first", "last"]:
return self.get_frame_names()
elif tag == "type":
return self.get_valid_types()
elif tag == "format":
return list(self.formats)
elif tag == "file_type":
return ["png", "jpeg", "tif", "exr"]
elif tag == "workers":
return common.get_workers(self)
elif tag == "ec2":
return self.get_ec2_instance_types() if self.parent.is_aws else []
else:
raise Exception(f"Invalid tag: {tag}")
def update_buttons(self, gb):
"""Enables buttons and dropdowns according to whether or not data is present.
Args:
gb (QtWidgets.QGroupBox): Group box for the tab.
"""
dlg = self.dlg
ignore = [dlg.dd_export_farm_workers, dlg.dd_export_farm_ec2]
if self.is_viewer_rift_hidden:
ignore.append(dlg.btn_export_data_rift_view)
if self.is_viewer_smr_hidden:
ignore.append(dlg.btn_export_data_smr_view)
if self.parent.is_aws:
dlg.btn_export_farm_terminate_cluster.setEnabled(
bool(len(common.get_aws_workers()))
)
common.update_buttons(self, gb, ignore)
if not self.is_viewer_rift_hidden:
has_fused = dep_util.check_image_existence(self.path_fused) != ""
dlg.btn_export_data_rift_view.setEnabled(has_fused)
def get_format(self, flags_from_data=True):
"""Gets format to populate the flagfile.
Args:
flags_from_data (dict[str, _]): Flags and their values, typically
pulled from the UI.
Returns:
str: Format in the UI that should populate the flagfile.
"""
format_text = self.dlg.dd_export_data_format.currentText()
formats = self.formats
if format_text:
return formats[format_text] if flags_from_data else format_text
else:
vals = formats.values() if flags_from_data else formats
return next(iter(vals))
def get_color_scale(self, dst_width, color_type):
if color_type == "background_color":
src_width = self.parent.bg_full_size_width
else:
src_width = self.parent.video_full_size_width
if not src_width or not dst_width:
return 1.0
return float(dst_width) / float(src_width)
def update_flags_from_data(self, flags):
"""Updates flags from the UI elements.
Args:
flags (dict[str, _]): Flags corresponding to the tab default binary.
"""
dlg = self.dlg
rig_fn = getattr(self, "path_rig_json", "")
flags["rig"] = rig_fn
flags["input_root"] = self.parent.path_project
flags["output_root"] = self.path_video
flags["log_dir"] = self.path_logs
flags["width"] = dlg.val_export_options_res.text()
flags["first"] = dlg.dd_export_data_first.currentText()
flags["last"] = dlg.dd_export_data_last.currentText()
flags["force_recompute"] = dlg.cb_export_recompute.isChecked()
color_type = dlg.dd_export_data_type.currentText()
if color_type == "background_color":
flags["color"] = self.path_bg_color
flags["disparity"] = self.path_bg_disparity
flags["color_type"] = "background_color"
flags["disparity_type"] = "background_disp"
else:
flags["color"] = self.path_video_color
flags["disparity"] = self.path_video_disparity
flags["color_type"] = "color"
flags["disparity_type"] = "disparity"
if self.is_farm and self.parent.is_aws:
flags["master"] = ""
flags["workers"] = ""
flags["cloud"] = "aws"
rig_bn = os.path.basename(flags["rig"])
flags["input_root"] = self.parent.ui_flags.project_root
flags["output_root"] = os.path.join(
self.parent.ui_flags.project_root, config.OUTPUT_ROOT_NAME
)
flags["rig"] = os.path.join(
self.parent.ui_flags.project_root, "rigs", rig_bn
)
flags["log_dir"] = os.path.join(flags["output_root"], "logs")
elif self.is_farm and self.parent.is_lan:
flags["master"] = self.parent.ui_flags.master
flags["workers"] = ",".join(dlg.dd_export_farm_workers.checkedItems())
flags["cloud"] = ""
flags["username"] = self.parent.ui_flags.username
flags["password"] = self.parent.ui_flags.password
else: # local
flags["master"] = ""
flags["workers"] = ""
flags["cloud"] = ""
flags["file_type"] = dlg.dd_export_data_file_type.currentText()
flags["format"] = self.get_format()
is6dof = flags["format"].startswith("6dof")
if is6dof:
flags["format"] = "6dof"
flags["run_convert_to_binary"] = (
is6dof and "Meshing" in dlg.dd_export_data_format.currentText()
)
flags["run_fusion"] = (
is6dof and "Striping" in dlg.dd_export_data_format.currentText()
)
flags["run_simple_mesh_renderer"] = not is6dof
flags["color_scale"] = self.get_color_scale(flags["width"], color_type)
if self.parent.is_aws:
create_flagfile = os.path.join(
self.path_flags, self.app_name_to_flagfile[self.app_aws_create]
)
if os.path.exists(create_flagfile):
create_flags = get_flags_from_flagfile(create_flagfile)
if "cluster_size" in create_flags:
dlg.spin_export_farm_num_workers.setValue(
int(create_flags["cluster_size"])
)
if "instance_type" in create_flags:
dlg.dd_export_farm_ec2.setCurrentText(create_flags["instance_type"])
def get_format_from_value(self, value):
"""Gets format to populate the UI from the flagfile. This is needed since the
format displayed on the UI and in the internal pipeline differ (for readability).
Args:
value (str): Value from the flagfile (i.e. eqrcolor).
Returns:
str: Format to be displayed in the UI (i.e. Equirect Color).
"""
for k, v in self.formats.items():
if v == value:
return k
return ""
def update_data_from_flags(self, flags):
"""Updates UI elements from the flags.
Args:
flags (dict[str, _]): Flags corresponding to the tab default binary.
"""
dlg = self.dlg
dropdowns = [
["color_type", dlg.dd_export_data_type],
["first", dlg.dd_export_data_first],
["last", dlg.dd_export_data_last],
["file_type", dlg.dd_export_data_file_type],
]
values = [["width", dlg.val_export_options_res]]
checkboxes = [["force_recompute", dlg.cb_export_recompute]]
common.update_data_from_flags(
self, flags, dropdowns=dropdowns, values=values, checkboxes=checkboxes
)
# Special case: format
if flags["format"] in self.formats.values():
val = self.get_format_from_value(flags["format"])
dep_util.update_qt_dropdown(dlg.dd_export_data_format, val)
def update_data_or_flags(
self, flagfile_fn, flagfile_from_data, switch_to_flag_tab=False
):
"""Updates the flagfile from the UI elements or vice versa.
Args:
flagfile_fn (str): Name of the flagfile.
flagfile_from_data (bool): Whether to load the flagfile from the data (True) or
vice versa (False).
switch_to_flag_tab (bool, optional): Whether or not to switch tabs after updating.
"""
common.update_data_or_flags(
self, flagfile_fn, flagfile_from_data, switch_to_flag_tab
)
def update_flagfile_edit(self, flagfile_fn, switch_to_flag_tab=True):
"""Updates the edit box for the flagfile.
Args:
flagfile_fn (str): Name of the flagfile.
switch_to_flag_tab (bool, optional): Whether or not to switch tabs after updating.
"""
common.update_flagfile_edit(self, flagfile_fn, switch_to_flag_tab)
def refresh_data(self):
"""Updates UI elements to be in sync with data on disk."""
common.refresh_data(self)
def setup_project(self, mkdirs=False):
"""Retrieves any missing flagfiles and sets the default flags for the tab.
Args:
mkdirs (bool, optional): Whether or not to make the directories where the
outputs are saved by default.
"""
common.setup_project(self, mkdirs)
common.setup_aws_config(self)
def refresh(self):
"""Resets the UI tab to its start state."""
self.setup_project()
def update_flagfile(self, flagfile_fn):
"""Updates the flagfile from UI elements.
Args:
flagfile_fn (str): Name of the flagfile.
"""
common.update_flagfile(self, flagfile_fn)
def save_flag_file(self):
"""Saves flagfile from the UI to disk."""
common.save_flag_file(self, self.flagfile_fn)
def retrieve_missing_flagfiles(self):
"""Copies the missing flagfiles to project for local modification."""
common.retrieve_missing_flagfiles(self)
def add_default_flags(self):
"""Retrieves the default flags to the local flagfile."""
common.add_default_flags(self)
def _ipc_call(self, ipc_fn):
"""Makes an IPC to the host. This is currently done through a Watchdog system
running on the host tied to specific files and us touching the corresponding
files from within the container.
Args:
ipc_fn (str): Path to the file being monitored on the host.
"""
if self.is_farm and self.parent.is_aws:
local_fused_dir = os.path.join(
config.DOCKER_OUTPUT_ROOT, image_type_paths["fused"]
)
if (
not os.path.exists(local_fused_dir)
or len(os.listdir(local_fused_dir)) == 0
):
remote_fused_dir = os.path.join(
self.parent.ui_flags.project_root,
config.OUTPUT_ROOT_NAME,
image_type_paths["fused"],
)
download(src=remote_fused_dir, dst=local_fused_dir, filters=["*"])
docker_ipc_path = os.path.join(config.DOCKER_IPC_ROOT, ipc_fn)
with open(docker_ipc_path, "w"):
os.utime(docker_ipc_path, None)
def activate_ipc(self, ipc_name):
"""Runs an application. This only works if the app is set up with a corresponding
IPC file that is monitored by the Watchdog (defined in config.py).
Args:
ipc_name (str): Name of the binary (must be a case in config.get_app_name).
Otherwise this is a no-op.
"""
app_name = config.get_app_name(ipc_name)
self.log_reader.log_notice(
f"Executing {app_name} in host. Check terminal for logs."
)
self._ipc_call(ipc_name)
def is_local_render(self):
"""Whether or not the current render is configured to be performed locally.
Returns:
bool: If the render will be performed locally.
"""
return (
(not self.parent.is_aws or not self.is_farm)
and not self.is_host_linux_gpu
and self.is_local_non_win
and self.get_format() != "6dof"
)
def run_process(self):
"""Runs the default binary associated with the tab."""
gb = self.dlg.gb_export_data
p_id = f"run_{self.tag}_{self.get_format()}"
if self.is_farm and self.parent.is_aws:
common.run_process_aws(self, gb, p_id=p_id)
elif self.is_local_render():
# Render non-6DoF through host call
self.activate_ipc(config.DOCKER_SMR_IPC)
else:
# Render
common.run_process(self, gb, p_id=p_id)
def cancel_process(self):
"""Stops a running process."""
common.cancel_process(self)
def reset_run_button_text(self):
self.dlg.btn_export_data_run.setText("Run")
def update_run_button_text(self):
"""Updates the text of the Run button depending on the existance of a process
running on the cloud
"""
common.update_run_button_text(self, self.dlg.btn_export_data_run)
def on_changed_flagfile_edit(self):
"""Callback event handler for flagfile edits."""
caller = sys._getframe().f_back.f_code.co_name
if caller == "main":
# Modified by typing on it
btn = self.dlg.btn_export_flagfile_save
if not btn.isEnabled():
btn.setEnabled(True)
def on_ec2_dashboard(self):
"""Callback event handler for clicking the "EC2 Dashboard" button."""
common.popup_ec2_dashboard_url(self)
def on_download_meshes(self):
"""Callback event handler for clicking the "Download Meshes" button."""
gb = self.dlg.gb_export_data
common.on_download_meshes(self, gb)
def get_logs(self):
"""Callback event handler for clicking the "Logs" button."""
common.popup_logs_locations(self)
def setup_flags(self):
"""Sets up the flags according to the corresponding flagfile."""
common.setup_flagfile_tab(self)
self.dlg.label_export_flagfile_tooltip.setToolTip(self.tooltip)
def update_viewer_buttons(self):
"""Only enable viewer buttons based on what is possible from the host, i.e.
RiftViewer will be disabled on any non-Windows computer.
"""
btn_smr_onscreen = self.dlg.btn_export_data_smr_view
btn_riftviewer = self.dlg.btn_export_data_rift_view
if self.get_format() != "6dof":
btn_smr_onscreen.setEnabled(False)
btn_riftviewer.setEnabled(False)
else:
if not self.is_viewer_smr_hidden:
btn_smr_onscreen.setEnabled(True)
if not self.is_viewer_rift_hidden:
btn_riftviewer.setEnabled(True)
def update_ec2_dropdown(self):
"""Displays only valid values in the dropdown for EC2 instance types."""
gb = self.dlg.gb_export_farm
dd = self.dlg.dd_export_farm_ec2
common.call_force_refreshing(self, common.populate_dropdown, self, gb, dd)
def update_file_type(self):
"""Enables/disables file type depending on the format"""
is6dof = self.get_format().startswith("6dof")
self.dlg.dd_export_data_file_type.setEnabled(not is6dof)
def update_label_resolution(self):
format = self.get_format()
if format.startswith("6dof"):
label = "per camera, max recommended 3072"
elif format.startswith("cube"):
label = "per cubeface"
else:
label = "equirect"
label = f"Output resolution ({label})"
self.dlg.label_export_options_res.setText(label)
def on_state_changed_recompute(self):
"""Callback event handler for clicking the "Re-compute" checkbox."""
common.on_state_changed_recompute(self)
def update_frame_range_dropdowns(self):
"""Updates ranges displayed in dropdowns per available files on disk."""
dlg = self.dlg
gb = dlg.gb_export_data
dds = [dlg.dd_export_data_first, dlg.dd_export_data_last]
for dd in dds:
common.call_force_refreshing(self, common.populate_dropdown, self, gb, dd)
def on_changed_dropdown(self, gb, dd):
"""Callback event handler for changed dropdown.
Args:
gb (QtWidgets.QGroupBox): Group box for the tab.
dd (QtWidgets.QComboBox): Dropdown UI element.
"""
common.on_changed_dropdown(self, gb, dd)
self.update_viewer_buttons()
if dd.objectName().endswith("_data_format"):
self.update_ec2_dropdown()
self.update_file_type()
self.update_label_resolution()
elif dd.objectName().endswith("_data_type"):
self.update_frame_range_dropdowns()
def on_changed_line_edit(self, gb, le):
"""Callback event handler for changed line edit.
Args:
gb (QtWidgets.QGroupBox): Group box for the tab.
le (QtWidgets.QLineEdit): Line edit UI element.
"""
common.on_changed_line_edit(self, gb, le)
def on_state_changed_alpha_blend(self):
"""Callback event handler for changed alpha value."""
if not self.is_refreshing_data:
self.update_flagfile(self.flagfile_fn)
def on_state_changed_farm(self, state):
"""Callback event handler for changed "AWS" checkbox.
Args:
state (str): Identifier of the callback state.
"""
common.on_state_changed_farm(self, state)
def on_terminate_cluster(self):
"""Terminates a running cluster."""
gb = self.dlg.gb_export_data
common.on_terminate_cluster(self, gb)
def on_changed_preview(self):
"""Callback event handler for changed image previews."""
common.on_changed_preview(self)
def populate_dropdowns(self, gb):
"""Populates the dropdowns in the tab.
Args:
gb (QtWidgets.QGroupBox): Group box for the tab.
"""
if gb == self.dlg.gb_export_data:
dds_pri = [self.dlg.dd_export_data_type]
else:
dds_pri = []
common.populate_dropdowns(self, gb, dds_pri)
def setup_data(self):
"""Sets up callbacks and initial UI element statuses."""
dlg = self.dlg
dlg.cb_export_recompute.setChecked(False)
dlg.label_export_farm_workers.setEnabled(False)
dlg.dd_export_farm_workers.setEnabled(False)
dlg.cb_export_alpha_blend.setChecked(True)
dlg.gb_export_farm.setEnabled(True)
dlg.btn_export_data_download_meshes.setEnabled(self.parent.is_aws)
callbacks = {
dlg.cb_export_recompute: self.on_state_changed_recompute,
dlg.btn_export_data_smr_view: (
lambda: self.activate_ipc(config.DOCKER_SMR_ONSCREEN_IPC)
),
dlg.btn_export_data_rift_view: (
lambda: self.activate_ipc(config.DOCKER_RIFT_VIEWER_IPC)
),
dlg.cb_export_alpha_blend: self.on_state_changed_alpha_blend,
dlg.btn_export_farm_terminate_cluster: self.on_terminate_cluster,
dlg.btn_export_farm_ec2_dashboard: self.on_ec2_dashboard,
dlg.btn_export_data_download_meshes: self.on_download_meshes,
dlg.gb_export_farm: self.on_state_changed_farm,
}
common.setup_data(self, callbacks)
def disable_tab_if_no_data(self):
"""Prevents navigation to the tab if the required data is not present."""
common.disable_tab_if_no_data(self, self.dlg.btn_export_data_run)