"""Tensorflow Ultrasound Beamforming Library.Initialize modules for registries."""importsysfrompathlibimportPosixPathimportnumpyasnp# Convert PosixPath objects to strings in sys.path# this is necessary due to weird TF bug when importingsys.path=[str(p)ifisinstance(p,PosixPath)elsepforpinsys.path]importtensorflowastf# noqa: E402from.dataloaderimportmake_dataloader# noqa: E402
[docs]defon_device_tf(func,inputs,device,return_numpy=False,**kwargs):"""Applies a Tensorflow function to inputs on a specified device. Args: func (function): Function to apply to the input data. inputs (ndarray): Input array. device (str): Device string, e.g. ``'cuda'``, ``'gpu'``, or ``'cpu'``. return_numpy (bool, optional): Whether to convert output data back to numpy. Defaults to False. **kwargs: Additional keyword arguments to be passed to the ``func``. Returns: tf.Tensor or ndarray: The output data. Raises: AssertionError: If ``func`` is not a function from the tensorflow library. Note: This function converts the ``inputs`` array to a tf.Tensor and moves it to the specified ``device``. It then applies the ``func`` function to the inputs and returns the output data. If the output is a dictionary, it extracts the first value from the dictionary. If ``return_numpy`` is True, it converts the output data back to a numpy array before returning. Example: .. code-block:: python import tensorflow as tf def square(x): return x**2 inputs = [1, 2, 3, 4, 5] device = "cuda" output = on_device_tf(square, inputs, device) """device=device.replace("cuda","gpu")withtf.device(device):outputs=func(inputs,**kwargs)ifreturn_numpy:ifnotisinstance(outputs,np.ndarray):outputs=outputs.numpy()returnoutputs