1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
use std::sync::{Arc, Mutex};
use std::collections::HashMap;
use tracing::info;
use ndarray::prelude::*;
use anyhow::{Context, Error, Result, Ok, bail, anyhow};
use common::message::*;
use common::geometry::Geometry;
use common::connection::Connection;
/// Rust interface for LibreDR client
#[derive(Debug)]
pub struct LibreDR {
connection: Connection,
}
impl LibreDR {
/// Construct `LibreDR` by connecting to LibreDR server\
/// Return `Error` if connection failed\
/// # Examples
/// ```
/// async {
/// let client_tcp = LibreDR::new(String::from("127.0.0.1:9001"), false, false).await?;
/// let client_unix = LibreDR::new(String::from("/var/run/libredr_client.sock"), true, false).await?;
/// }
/// ```
pub async fn new(connect: String, unix: bool, tls: bool) -> Result<Self> {
info!("LibreDR::new: Connecting to server {connect}");
let mut config: HashMap<String, String> = HashMap::new();
config.insert(String::from("connect"), connect.to_owned());
config.insert(String::from("unix"), String::from(if unix { "true" } else { "false" }));
config.insert(String::from("tls"), String::from(if tls { "true" } else { "false" }));
let connection = Connection::from_config(&config).await
.with_context(|| format!("Failed to connect to server {connect}"))?;
Ok(LibreDR {
connection,
})
}
/// Receive messages, response `RequestData` task, until receive a different task
async fn try_recv_msg_response_data(&mut self, data_cache: &DataCache) -> Result<Message> {
loop {
let msg_response = self.connection.recv_msg().await?;
info!("LibreDR::try_recv_msg_response_data: msg_response {msg_response}");
if let Message::RequestData(hash) = msg_response {
let data = {
let data_cache = data_cache.lock().expect("No task should panic");
let entry = data_cache.get(&hash).ok_or(format!("Client: ray_tracing_forward: unexpected hash {hash}"));
entry.map(|entry| {
entry.1.to_owned()
})
};
let msg_response = Message::ResponseData(Box::new(data.to_owned()));
self.connection.send_msg(&msg_response).await?;
data.map_err(Error::msg)?;
} else {
break Ok(msg_response);
}
}
}
/// Create a [`RequestRayTracingForward`] task and wait for response
///
/// # Arguments
/// * `ray` - ray parameters
/// * if `camera_space` is `false` 18 * `image_shape`
/// * including ray position 9 * `image_shape`
/// * including ray direction 9 * `image_shape`
/// * if `camera_space` is `true`, add another (1 + 14) channels
/// * including ray depth 1 * `image_shape` (if depth <= 0, treat as hit miss)
/// * including ray material 14 * `image_shape`
/// * `texture` - (3 + 3 + 3 + 1 + 3 + 1) * `texture_resolution` * `texture_resolution` (must be square image)
/// * including normal + diffuse + specular + roughness + intensity + window
/// * `envmap` - 3 * 6 * `envmap_resolution` * `envmap_resolution`
/// * (must be box unwrapped 6 square images)
/// * `sample_per_pixel` - `sample_per_pixel_forward`, `sample_per_pixel_backward`
/// * `max_bounce` - `max_bounce_forward`, `max_bounce_backward`, `max_bounce_low_discrepancy`, `skip_bounce`
/// * `switches` - tuple of 4 switches to determine hit miss and reflection behavior
/// * render::MISS_* - determine how to deal with ray hit miss
/// * [`common::render::MISS_NONE`]
/// * [`common::render::MISS_ENVMAP`]
/// * render::REFLECTION_NORMAL_* - determine how to get surface normal
/// * [`common::render::REFLECTION_NORMAL_FACE`]
/// * [`common::render::REFLECTION_NORMAL_VERTEX`]
/// * [`common::render::REFLECTION_NORMAL_TEXTURE`]
/// * render::REFLECTION_DIFFUSE_* - determine diffuse reflection model
/// * [`common::render::REFLECTION_DIFFUSE_NONE`]
/// * [`common::render::REFLECTION_DIFFUSE_LAMBERTIAN`]
/// * render::REFLECTION_SPECULAR_* - determine specular reflection model
/// * [`common::render::REFLECTION_SPECULAR_NONE`]
/// * [`common::render::REFLECTION_SPECULAR_PHONG`]
/// * [`common::render::REFLECTION_SPECULAR_BLINN_PHONG`]
/// * [`common::render::REFLECTION_SPECULAR_TORRANCE_SPARROW_PHONG`]
/// * [`common::render::REFLECTION_SPECULAR_TORRANCE_SPARROW_BLINN_PHONG`]
/// * [`common::render::REFLECTION_SPECULAR_TORRANCE_SPARROW_BECKMANN`]
/// * `clip_near` - clip near distance for camera
/// * `clip_near` can be a single float number (same for all bounces),
/// * or tuple of 3 float numbers (first bounce, second bounce, and other bounces)
/// * `camera_space` - if `true`, the first bounce uses the depth and material given by the ray
/// * `requires_grad` - if `true`, worker will save intermediate data, the next task must be `ray_tracing_backward`
/// * `srand` - random seed
/// * if srand >= 0, the same random seed is used for every pixel
/// * if srand < 0, use different seed for each pixel
/// * `low_discrepancy` - (optional) start id of Halton low discrepancy sequence.
/// * The default value is the same as `sample_per_pixel_forward`.
/// * if combine multiple rendered images to reduce noise, this value can be set to: \
/// 1 * `sample_per_pixel_forward`, 2 * `sample_per_pixel_forward`, 3 * `sample_per_pixel_forward`, ...
///
/// # Return
/// Return shape will be,
/// * if `camera_space` is `true`
/// * render image 3 * `image_shape`
/// * if `camera_space` is `false`, add another
/// * ray texture coordinate 2 * `image_shape`
/// * ray depth (Euclidean distance) 1 * `image_shape`
/// * ray normal 3 * `image_shape`
#[allow(clippy::too_many_arguments)]
pub async fn ray_tracing_forward(&mut self,
geometry: &Geometry,
geometry_data_cache: &DataCache,
ray: ArrayD<f32>,
texture: Array3<f32>,
envmap: Array4<f32>,
sample_per_pixel: (usize, usize),
max_bounce: (usize, usize, usize, usize),
switches: (u8, u8, u8, u8),
clip_near: (f32, f32, f32),
camera_space: bool,
requires_grad: bool,
srand: i32,
low_discrepancy: u32) -> Result<ArrayD<f32>> {
let input_shape = ray.shape().to_owned();
assert!(input_shape.len() > 1, "ray_tracing_forward: ray should be at least 1D");
let ray_channels_input = if camera_space { 33 } else { 18 };
assert_eq!(input_shape[0], ray_channels_input,
"ray_tracing_forward: ray channel: {}, expected: {}", input_shape[0], ray_channels_input);
let ray = ray.to_shape([ray_channels_input, input_shape[1..].iter().product()])?.into_owned();
let mut data_cache_content = hashbrown::HashMap::new();
{
let geometry_data_cache = geometry_data_cache.lock().expect("No task should panic");
data_cache_content.extend(geometry_data_cache.iter().map(|(k, v)| {
(k.to_owned(), v.to_owned())
}));
}
let ray_data = Data::RayData(ray);
let ray_data_hash = ray_data.hash();
data_cache_content.insert(ray_data_hash.to_owned(), (0, ray_data));
let material_data = Data::MaterialData(texture, envmap);
let material_data_hash = material_data.hash();
data_cache_content.insert(material_data_hash.to_owned(), (0, material_data));
let data_cache = Arc::new(Mutex::new(data_cache_content));
let request = RequestRayTracingForward {
geometry: geometry.to_owned(),
ray: ray_data_hash,
material: material_data_hash,
sample_per_pixel,
max_bounce,
switches,
clip_near,
camera_space,
requires_grad,
srand,
low_discrepancy
};
let msg_request = Message::RequestTask(RequestTask::RequestRayTracingForward(Box::new(request)));
info!("LibreDR::ray_tracing_forward: msg_request {msg_request}");
self.connection.send_msg(&msg_request).await?;
let msg_response = self.try_recv_msg_response_data(&data_cache).await?;
let Message::ResponseTask(response_task) = msg_response else {
bail!("Unexpected response {msg_response}");
};
let response_task = response_task.map_err(Error::msg)?;
let ResponseTask::ResponseRayTracingForward(response) = response_task else {
bail!("Unexpected response {response_task}");
};
let mut output_shape = input_shape.to_owned();
output_shape[0] = if camera_space { 3 } else { 9 };
let response = response.render.to_shape(output_shape)?.into_owned();
Ok(response)
}
/// Create a [`RequestRayTracingBackward`] task and wait for response.
///
/// Must be called consecutive to a [`RequestRayTracingForward`] task with `requires_grad` set to `true`. \
/// To create multiple [`RequestRayTracingForward`] tasks and backward together, multiple client connections are
/// required.
///
/// # Arguments
/// * `d_ray` - gradient of image 3 * `image_shape` (must ensure same `image_shape` as [`RequestRayTracingForward`])
///
/// # Return
/// Return shape will be,
/// * if `camera_space` is `false` for [`RequestRayTracingForward`] task
/// * 1st return value (3 + 3 + 3 + 1 + 3 + 1) * `texture_resolution` * `texture_resolution`
/// * (same `texture_resolution` as [`RequestRayTracingForward`])
/// * including d_normal + d_diffuse + d_specular + d_roughness + d_intensity + d_window
/// * 2nd return value 3 * 6 * `envmap_resolution` * `envmap_resolution`
/// * (same `envmap_resolution` as [`RequestRayTracingForward`])
/// * including d_envmap
/// * if `camera_space` is `true` for [`RequestRayTracingForward`] task, add another
/// * 3rd return value 14 * `image_shape` (same shape as [`RequestRayTracingForward`])
/// * including d_ray_texture
#[allow(clippy::too_many_arguments)]
pub async fn ray_tracing_backward(&mut self, d_ray: ArrayD<f32>) ->
Result<(Array3<f32>, Array4<f32>, Option<ArrayD<f32>>)> {
let input_shape = d_ray.shape().to_owned();
assert!(input_shape.len() > 1, "ray_tracing_backward: d_ray should be at least 1D");
assert_eq!(input_shape[0], 3,
"ray_tracing_backward: d_ray channel: {}, expected: {}", input_shape[0], 3);
let d_ray = d_ray.to_shape([3, input_shape[1..].iter().product()])?.into_owned();
let request = RequestRayTracingBackward { d_ray, };
let msg_request = Message::RequestTask(RequestTask::RequestRayTracingBackward(Box::new(request)));
info!("LibreDR::ray_tracing_backward: msg_request {msg_request}");
self.connection.send_msg(&msg_request).await?;
let msg_response = self.connection.recv_msg().await?;
info!("LibreDR::ray_tracing_backward: msg_response {msg_response}");
let Message::ResponseTask(response_task) = msg_response else {
bail!("Unexpected response {msg_response}");
};
let response_task = response_task.map_err(Error::msg)?;
let ResponseTask::ResponseRayTracingBackward(response) = response_task else {
bail!("Unexpected response {response_task}");
};
let d_ray_texture = response.d_ray_texture.map(|d_ray_texture| {
let mut output_shape = input_shape.to_owned();
output_shape[0] = 14;
Ok(d_ray_texture.to_shape(output_shape)?.into_owned())
}).transpose()?;
Ok((
response.d_texture.ok_or(anyhow!("LibreDR::ray_tracing_backward: None d_texture returned from server"))?,
response.d_envmap.ok_or(anyhow!("LibreDR::ray_tracing_backward: None d_envmap returned from server"))?,
d_ray_texture))
}
/// Send [`Message::Close`] to server to close cleanly
pub async fn close(&mut self) -> Result<()> {
self.connection.send_msg(&Message::Close()).await
}
}