import os
import numpy as np
from blenderproc.python.modules.main.Module import Module
from blenderproc.python.utility.Utility import Utility, resolve_path
from blenderproc.python.writer.WriterUtility import WriterUtility
[docs]class WriterInterface(Module):
"""
Parent class for all other writers classes, it had the functionality to return objects attributes and write \
them to file and to load and process post processing modules
**Configuration**:
.. list-table::
:widths: 25 100 10
:header-rows: 1
* - Parameter
- Description
- Type
* - postprocessing_modules
- A dict of list of postprocessing modules. The key in the dict specifies the output to which the
postprocessing modules should be applied. Every postprocessing module has to have a run function which
takes in the raw data and returns the processed data.
- dict
* - destination_frame
- Used to transform point to blender coordinate frame. Default: ["X", "Y", "Z"]
- list
* - attributes_to_write
- A list of attribute names that should written to file. The next table lists all attributes that can be
used here.
- list
* - output_file_prefix
- The prefix of the file that should be created.
- string
* - output_key
- The key which should be used for storing the output in a merged file.
- string
* - write_alpha_channel
- If true, the alpha channel will be written to file. Default: False.
- bool
"""
def __init__(self, config):
Module.__init__(self, config)
self.postprocessing_modules_per_output = {}
module_configs = config.get_raw_dict("postprocessing_modules", {})
for output_key in module_configs:
self.postprocessing_modules_per_output[output_key] = Utility.initialize_modules(module_configs[output_key])
self.name_to_id = {}
self.destination_frame = self.config.get_list("destination_frame", ["X", "Y", "Z"])
self.write_alpha_channel = self.config.get_bool("write_alpha_channel", False)
[docs] def write_attributes_to_file(self, item_writer, items, default_file_prefix, default_output_key, default_attributes, version="1.0.0"):
""" Writes the state of the given items to a file with the configured prefix.
This method also registers the corresponding output.
:param item_writer: The item writer object to use. Type: object.
:param items: The list of items. Type: list.
:param default_file_prefix: The default file name prefix to use. Type: string.
:param default_output_key: The default output key to use. Type: string.
:param default_attributes: The default attributes to write, if no attributes are specified in the config. Type: list.
:param version: The version to use when registering the output. Type: string.
"""
if self._avoid_output:
print("Avoid output is on, no output produced!")
return
file_prefix = self.config.get_string("output_file_prefix", default_file_prefix)
path_prefix = os.path.join(self._determine_output_dir(), file_prefix)
item_writer.write_items_to_file(path_prefix, items, self.config.get_list("attributes_to_write", default_attributes), world_frame_change=self.destination_frame)
Utility.register_output(self._determine_output_dir(), file_prefix, self.config.get_string("output_key", default_output_key), ".npy", version)
[docs] def _apply_postprocessing(self, output_key: str, data: np.ndarray, version: str):
"""
Applies all postprocessing modules registered for this output type.
:param output_key: The key of the output type. Type: string
:param data: The numpy data.
:param version: The version number original data.
:return: The modified numpy data after doing the postprocessing
"""
if output_key in self.postprocessing_modules_per_output:
for module in self.postprocessing_modules_per_output[output_key]:
data, new_key, new_version = module.run(data, output_key, version)
else:
new_key = output_key
new_version = version
return data, new_key, new_version
[docs] def _load_and_postprocess(self, file_path, key, version = "1.0.0"):
"""
Loads an image and post process it.
:param file_path: Image path. Type: string.
:param key: The image's key with regards to the hdf5 file. Type: string.
:param version: The version number original data. Type: String. Default: 1.0.0.
:return: The post-processed image that was loaded using the file path.
"""
data = WriterUtility.load_output_file(resolve_path(file_path), self.write_alpha_channel, remove=False)
data, new_key, new_version = self._apply_postprocessing(key, data, version)
if isinstance(data, np.ndarray):
print("Key: " + key + " - shape: " + str(data.shape) + " - dtype: " + str(data.dtype) + " - path: " + file_path)
else:
print("Key: " + key + " - path: " + file_path)
return data, new_key, new_version