2fceee09c788c700d77ceec6eb46864a1b41bd11
[blender.git] / intern / cycles / render / denoising.cpp
1 /*
2  * Copyright 2011-2018 Blender Foundation
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16
17 #include "render/denoising.h"
18
19 #include "kernel/filter/filter_defines.h"
20
21 #include "util/util_foreach.h"
22 #include "util/util_map.h"
23 #include "util/util_system.h"
24 #include "util/util_time.h"
25
26 #include <OpenImageIO/filesystem.h>
27
28 CCL_NAMESPACE_BEGIN
29
30 /* Utility Functions */
31
32 static void print_progress(int num, int total, int frame, int num_frames)
33 {
34         const char *label = "Denoise Frame ";
35         int cols = system_console_width();
36
37         cols -= strlen(label);
38
39         int len = 1;
40         for(int x = total; x > 9; x /= 10) {
41                 len++;
42         }
43
44         int bars = cols - 2*len - 6;
45
46         printf("\r%s", label);
47
48         if(num_frames > 1) {
49                 int frame_len = 1;
50                 for(int x = num_frames - 1; x > 9; x /= 10) {
51                         frame_len++;
52                 }
53                 bars -= frame_len + 2;
54                 printf("%*d ", frame_len, frame);
55         }
56
57         int v = int(float(num)*bars/total);
58         printf("[");
59         for(int i = 0; i < v; i++) {
60                 printf("=");
61         }
62         if(v < bars) {
63                 printf(">");
64         }
65         for(int i = v+1; i < bars; i++) {
66                 printf(" ");
67         }
68         printf(string_printf("] %%%dd / %d", len, total).c_str(), num);
69         fflush(stdout);
70 }
71
72 /* Splits in at its last dot, setting suffix to the part after the dot and in to the part before it.
73  * Returns whether a dot was found. */
74 static bool split_last_dot(string &in, string &suffix)
75 {
76         size_t pos = in.rfind(".");
77         if(pos == string::npos) {
78                 return false;
79         }
80         suffix = in.substr(pos+1);
81         in = in.substr(0, pos);
82         return true;
83 }
84
85 /* Separate channel names as generated by Blender.
86  * If views is true:
87  *   Inputs are expected in the form RenderLayer.Pass.View.Channel, sets renderlayer to "RenderLayer.View"
88  * Otherwise:
89  *   Inputs are expected in the form RenderLayer.Pass.Channel */
90 static bool parse_channel_name(string name, string &renderlayer, string &pass, string &channel, bool multiview_channels)
91 {
92         if(!split_last_dot(name, channel)) {
93                 return false;
94         }
95         string view;
96         if(multiview_channels && !split_last_dot(name, view)) {
97                 return false;
98         }
99         if(!split_last_dot(name, pass)) {
100                 return false;
101         }
102         renderlayer = name;
103
104         if(multiview_channels) {
105                 renderlayer += "." + view;
106         }
107
108         return true;
109 }
110
111 /* Channel Mapping */
112
113 struct ChannelMapping {
114         int channel;
115         string name;
116 };
117
118 static void fill_mapping(vector<ChannelMapping> &map, int pos, string name, string channels)
119 {
120         for(const char *chan = channels.c_str(); *chan; chan++) {
121                 map.push_back({pos++, name + "." + *chan});
122         }
123 }
124
125 static const int INPUT_NUM_CHANNELS = 15;
126 static vector<ChannelMapping> input_channels()
127 {
128         vector<ChannelMapping> map;
129         fill_mapping(map, 0, "Denoising Depth", "Z");
130         fill_mapping(map, 1, "Denoising Normal", "XYZ");
131         fill_mapping(map, 4, "Denoising Shadowing", "X");
132         fill_mapping(map, 5, "Denoising Albedo", "RGB");
133         fill_mapping(map, 8, "Noisy Image", "RGB");
134         fill_mapping(map, 11, "Denoising Variance", "RGB");
135         fill_mapping(map, 14, "Denoising Intensity", "X");
136         return map;
137 }
138
139 static const int OUTPUT_NUM_CHANNELS = 3;
140 static vector<ChannelMapping> output_channels()
141 {
142         vector<ChannelMapping> map;
143         fill_mapping(map, 0, "Combined", "RGB");
144         return map;
145 }
146
147 /* Renderlayer Handling */
148
149 bool DenoiseImageLayer::detect_denoising_channels()
150 {
151         /* Map device input to image channels. */
152         input_to_image_channel.clear();
153         input_to_image_channel.resize(INPUT_NUM_CHANNELS, -1);
154
155         foreach(const ChannelMapping& mapping, input_channels()) {
156                 vector<string>::iterator i = find(channels.begin(), channels.end(), mapping.name);
157                 if(i == channels.end()) {
158                         return false;
159                 }
160
161                 size_t input_channel = mapping.channel;
162                 size_t layer_channel = i - channels.begin();
163                 input_to_image_channel[input_channel] = layer_to_image_channel[layer_channel];
164         }
165
166         /* Map device output to image channels. */
167         output_to_image_channel.clear();
168         output_to_image_channel.resize(OUTPUT_NUM_CHANNELS, -1);
169
170         foreach(const ChannelMapping& mapping, output_channels()) {
171                 vector<string>::iterator i = find(channels.begin(), channels.end(), mapping.name);
172                 if(i == channels.end()) {
173                         return false;
174                 }
175
176                 size_t output_channel = mapping.channel;
177                 size_t layer_channel = i - channels.begin();
178                 output_to_image_channel[output_channel] = layer_to_image_channel[layer_channel];
179         }
180
181         /* Check that all buffer channels are correctly set. */
182         for(int i = 0; i < INPUT_NUM_CHANNELS; i++) {
183                 assert(input_to_image_channel[i] >= 0);
184         }
185         for(int i = 0; i < OUTPUT_NUM_CHANNELS; i++) {
186                 assert(output_to_image_channel[i] >= 0);
187         }
188
189         return true;
190 }
191
192 bool DenoiseImageLayer::match_channels(int neighbor,
193                                        const std::vector<string> &channelnames,
194                                        const std::vector<string> &neighbor_channelnames)
195 {
196         neighbor_input_to_image_channel.resize(neighbor + 1);
197         vector<int>& mapping = neighbor_input_to_image_channel[neighbor];
198
199         assert(mapping.size() == 0);
200         mapping.resize(input_to_image_channel.size(), -1);
201
202         for(int i = 0; i < input_to_image_channel.size(); i++) {
203                 const string& channel = channelnames[input_to_image_channel[i]];
204                 std::vector<string>::const_iterator frame_channel = find(neighbor_channelnames.begin(), neighbor_channelnames.end(), channel);
205
206                 if(frame_channel == neighbor_channelnames.end()) {
207                         return false;
208                 }
209
210                 mapping[i] = frame_channel - neighbor_channelnames.begin();
211         }
212
213         return true;
214 }
215
216 /* Denoise Task */
217
218 DenoiseTask::DenoiseTask(Device *device, Denoiser *denoiser, int frame, const vector<int>& neighbor_frames)
219 : denoiser(denoiser),
220   device(device),
221   frame(frame),
222   neighbor_frames(neighbor_frames),
223   current_layer(0),
224   input_pixels(device, "filter input buffer", MEM_READ_ONLY),
225   num_tiles(0)
226 {
227         image.samples = denoiser->samples_override;
228 }
229
230 DenoiseTask::~DenoiseTask()
231 {
232         free();
233 }
234
235 /* Device callbacks */
236
237 bool DenoiseTask::acquire_tile(Device *device, Device *tile_device, RenderTile &tile)
238 {
239         thread_scoped_lock tile_lock(tiles_mutex);
240
241         if(tiles.empty()) {
242                 return false;
243         }
244
245         tile = tiles.front();
246         tiles.pop_front();
247
248         device->map_tile(tile_device, tile);
249
250         print_progress(num_tiles - tiles.size(), num_tiles, frame, denoiser->num_frames);
251
252         return true;
253 }
254
255 /* Mapping tiles is required for regular rendering since each tile has its separate memory
256  * which may be allocated on a different device.
257  * For standalone denoising, there is a single memory that is present on all devices, so the only
258  * thing that needs to be done here is to specify the surrounding tile geometry.
259  *
260  * However, since there is only one large memory, the denoised result has to be written to
261  * a different buffer to avoid having to copy an entire horizontal slice of the image. */
262 void DenoiseTask::map_neighboring_tiles(RenderTile *tiles, Device *tile_device)
263 {
264         for(int i = 0; i < 9; i++) {
265                 if(i == 4) {
266                         continue;
267                 }
268
269                 int dx = (i%3)-1;
270                 int dy = (i/3)-1;
271                 tiles[i].x = clamp(tiles[4].x +  dx   *denoiser->tile_size.x, 0, image.width);
272                 tiles[i].w = clamp(tiles[4].x + (dx+1)*denoiser->tile_size.x, 0, image.width) - tiles[i].x;
273                 tiles[i].y = clamp(tiles[4].y +  dy   *denoiser->tile_size.y, 0, image.height);
274                 tiles[i].h = clamp(tiles[4].y + (dy+1)*denoiser->tile_size.y, 0, image.height) - tiles[i].y;
275
276                 tiles[i].buffer = tiles[4].buffer;
277                 tiles[i].offset = tiles[4].offset;
278                 tiles[i].stride = image.width;
279         }
280
281         device_vector<float> *output_mem = new device_vector<float>(tile_device, "denoising_output", MEM_READ_WRITE);
282         output_mem->alloc(OUTPUT_NUM_CHANNELS*tiles[4].w*tiles[4].h);
283         output_mem->zero_to_device();
284
285         tiles[9] = tiles[4];
286         tiles[9].buffer = output_mem->device_pointer;
287         tiles[9].stride = tiles[9].w;
288         tiles[9].offset -= tiles[9].x + tiles[9].y*tiles[9].stride;
289
290         thread_scoped_lock output_lock(output_mutex);
291         assert(output_pixels.count(tiles[4].tile_index) == 0);
292         output_pixels[tiles[9].tile_index] = output_mem;
293 }
294
295 void DenoiseTask::unmap_neighboring_tiles(RenderTile *tiles)
296 {
297         thread_scoped_lock output_lock(output_mutex);
298         assert(output_pixels.count(tiles[4].tile_index) == 1);
299         device_vector<float> *output_mem = output_pixels[tiles[9].tile_index];
300         output_pixels.erase(tiles[4].tile_index);
301         output_lock.unlock();
302
303         output_mem->copy_from_device(0, OUTPUT_NUM_CHANNELS*tiles[9].w, tiles[9].h);
304
305         float *result = output_mem->data();
306         float *out = &image.pixels[image.num_channels*(tiles[9].y*image.width + tiles[9].x)];
307
308         const DenoiseImageLayer& layer = image.layers[current_layer];
309         const int *output_to_image_channel = layer.output_to_image_channel.data();
310
311         for(int y = 0; y < tiles[9].h; y++) {
312                 for(int x = 0; x < tiles[9].w; x++, result += OUTPUT_NUM_CHANNELS) {
313                         for(int i = 0; i < OUTPUT_NUM_CHANNELS; i++) {
314                                 out[image.num_channels*x + output_to_image_channel[i]] = result[i];
315                         }
316                 }
317                 out += image.num_channels * image.width;
318         }
319
320         output_mem->free();
321         delete output_mem;
322 }
323
324 void DenoiseTask::release_tile()
325 {
326 }
327
328 bool DenoiseTask::get_cancel()
329 {
330         return false;
331 }
332
333 void DenoiseTask::create_task(DeviceTask& task)
334 {
335         /* Callback functions. */
336         task.acquire_tile = function_bind(&DenoiseTask::acquire_tile, this, device, _1, _2);
337         task.map_neighbor_tiles = function_bind(&DenoiseTask::map_neighboring_tiles, this, _1, _2);
338         task.unmap_neighbor_tiles = function_bind(&DenoiseTask::unmap_neighboring_tiles, this, _1);
339         task.release_tile = function_bind(&DenoiseTask::release_tile, this);
340         task.get_cancel = function_bind(&DenoiseTask::get_cancel, this);
341
342         /* Denoising parameters. */
343         task.denoising = denoiser->params;
344         task.denoising_do_filter = true;
345         task.denoising_write_passes = false;
346         task.denoising_from_render = false;
347
348         task.denoising_frames.resize(neighbor_frames.size());
349         for(int i = 0; i < neighbor_frames.size(); i++) {
350                 task.denoising_frames[i] = neighbor_frames[i] - frame;
351         }
352
353         /* Buffer parameters. */
354         task.pass_stride = INPUT_NUM_CHANNELS;
355         task.target_pass_stride = OUTPUT_NUM_CHANNELS;
356         task.pass_denoising_data = 0;
357         task.pass_denoising_clean = -1;
358         task.frame_stride = image.width * image.height * INPUT_NUM_CHANNELS;
359
360         /* Create tiles. */
361         thread_scoped_lock tile_lock(tiles_mutex);
362         thread_scoped_lock output_lock(output_mutex);
363
364         tiles.clear();
365         assert(output_pixels.empty());
366         output_pixels.clear();
367
368         int tiles_x = divide_up(image.width, denoiser->tile_size.x);
369         int tiles_y = divide_up(image.height, denoiser->tile_size.y);
370
371         for(int ty = 0; ty < tiles_y; ty++) {
372                 for(int tx = 0; tx < tiles_x; tx++) {
373                         RenderTile tile;
374                         tile.x = tx * denoiser->tile_size.x;
375                         tile.y = ty * denoiser->tile_size.y;
376                         tile.w = min(image.width - tile.x, denoiser->tile_size.x);
377                         tile.h = min(image.height - tile.y, denoiser->tile_size.y);
378                         tile.start_sample = 0;
379                         tile.num_samples = image.layers[current_layer].samples;
380                         tile.sample = 0;
381                         tile.offset = 0;
382                         tile.stride = image.width;
383                         tile.tile_index = ty*tiles_x + tx;
384                         tile.task = RenderTile::DENOISE;
385                         tile.buffers = NULL;
386                         tile.buffer = input_pixels.device_pointer;
387                         tiles.push_back(tile);
388                 }
389         }
390
391         num_tiles = tiles.size();
392 }
393
394 /* Denoiser Operations */
395
396 bool DenoiseTask::load_input_pixels(int layer)
397 {
398         int w = image.width;
399         int h = image.height;
400         int num_pixels = image.width * image.height;
401         int frame_stride = num_pixels * INPUT_NUM_CHANNELS;
402
403         /* Load center image */
404         DenoiseImageLayer& image_layer = image.layers[layer];
405
406         float *buffer_data = input_pixels.data();
407         image.read_pixels(image_layer, buffer_data);
408         buffer_data += frame_stride;
409
410         /* Load neighbor images */
411         for(int i = 0; i < image.in_neighbors.size(); i++) {
412                 if(!image.read_neighbor_pixels(i, image_layer, buffer_data)) {
413                         error = "Failed to read neighbor frame pixels";
414                         return false;
415                 }
416                 buffer_data += frame_stride;
417         }
418
419         /* Preprocess */
420         buffer_data = input_pixels.data();
421         for(int neighbor = 0; neighbor < image.in_neighbors.size() + 1; neighbor++) {
422                 /* Clamp */
423                 if(denoiser->params.clamp_input) {
424                         for(int i = 0; i < num_pixels*INPUT_NUM_CHANNELS; i++) {
425                                 buffer_data[i] = clamp(buffer_data[i], -1e8f, 1e8f);
426                         }
427                 }
428
429                 /* Box blur */
430                 int r = 5 * denoiser->params.radius;
431                 float *data = buffer_data + 14;
432                 array<float> temp(num_pixels);
433
434                 for(int y = 0; y < h; y++) {
435                         for(int x = 0; x < w; x++) {
436                                 int n = 0;
437                                 float sum = 0.0f;
438                                 for(int dx = max(x - r, 0); dx < min(x + r + 1, w); dx++, n++) {
439                                         sum += data[INPUT_NUM_CHANNELS * (y * w + dx)];
440                                 }
441                                 temp[y * w + x] = sum/n;
442                         }
443                 }
444
445                 for(int y = 0; y < h; y++) {
446                         for(int x = 0; x < w; x++) {
447                                 int n = 0;
448                                 float sum = 0.0f;
449
450                                 for(int dy = max(y - r, 0); dy < min(y + r + 1, h); dy++, n++) {
451                                         sum += temp[dy * w + x];
452                                 }
453
454                                 data[INPUT_NUM_CHANNELS * (y * w + x)] = sum/n;
455                         }
456                 }
457
458                 buffer_data += frame_stride;
459         }
460
461         /* Copy to device */
462         input_pixels.copy_to_device();
463
464         return true;
465 }
466
467 /* Task stages */
468
469 bool DenoiseTask::load()
470 {
471         string center_filepath = denoiser->input[frame];
472         if(!image.load(center_filepath, error)) {
473                 return false;
474         }
475
476         if(!image.load_neighbors(denoiser->input, neighbor_frames, error)) {
477                 return false;
478         }
479
480         if(image.layers.empty()) {
481                 error = "No image layers found to denoise in " + center_filepath;
482                 return false;
483         }
484
485         /* Allocate device buffer. */
486         int num_frames = image.in_neighbors.size() + 1;
487         input_pixels.alloc(image.width * INPUT_NUM_CHANNELS, image.height * num_frames);
488         input_pixels.zero_to_device();
489
490         /* Read pixels for first layer. */
491         current_layer = 0;
492         if(!load_input_pixels(current_layer)) {
493                 return false;
494         }
495
496         return true;
497 }
498
499 bool DenoiseTask::exec()
500 {
501         for(current_layer = 0; current_layer < image.layers.size(); current_layer++) {
502                 /* Read pixels for secondary layers, first was already loaded. */
503                 if(current_layer > 0) {
504                         if(!load_input_pixels(current_layer)) {
505                                 return false;
506                         }
507                 }
508
509                 /* Run task on device. */
510                 DeviceTask task(DeviceTask::RENDER);
511                 create_task(task);
512                 device->task_add(task);
513                 device->task_wait();
514
515                 printf("\n");
516         }
517
518         return true;
519 }
520
521 bool DenoiseTask::save()
522 {
523         bool ok = image.save_output(denoiser->output[frame], error);
524         free();
525         return ok;
526 }
527
528 void DenoiseTask::free()
529 {
530         image.free();
531         input_pixels.free();
532         assert(output_pixels.empty());
533 }
534
535 /* Denoise Image Storage */
536
537 DenoiseImage::DenoiseImage()
538 {
539         width = 0;
540         height = 0;
541         num_channels = 0;
542         samples = 0;
543 }
544
545 DenoiseImage::~DenoiseImage()
546 {
547         free();
548 }
549
550 void DenoiseImage::close_input()
551 {
552         in_neighbors.clear();
553 }
554
555 void DenoiseImage::free()
556 {
557         close_input();
558         pixels.clear();
559 }
560
561 bool DenoiseImage::parse_channels(const ImageSpec &in_spec, string& error)
562 {
563         const std::vector<string> &channels = in_spec.channelnames;
564         const ParamValue *multiview = in_spec.find_attribute("multiView");
565         const bool multiview_channels = (multiview &&
566                                          multiview->type().basetype == TypeDesc::STRING &&
567                                          multiview->type().arraylen >= 2);
568
569         layers.clear();
570
571         /* Loop over all the channels in the file, parse their name and sort them
572          * by RenderLayer.
573          * Channels that can't be parsed are directly passed through to the output. */
574         map<string, DenoiseImageLayer> file_layers;
575         for(int i = 0; i < channels.size(); i++) {
576                 string layer, pass, channel;
577                 if(parse_channel_name(channels[i], layer, pass, channel, multiview_channels)) {
578                         file_layers[layer].channels.push_back(pass + "." + channel);
579                         file_layers[layer].layer_to_image_channel.push_back(i);
580                 }
581         }
582
583         /* Loop over all detected RenderLayers, check whether they contain a full set of input channels.
584          * Any channels that won't be processed internally are also passed through. */
585         for(map<string, DenoiseImageLayer>::iterator i = file_layers.begin(); i != file_layers.end(); ++i) {
586                 const string& name = i->first;
587                 DenoiseImageLayer& layer = i->second;
588
589                 /* Check for full pass set. */
590                 if(!layer.detect_denoising_channels()) {
591                         continue;
592                 }
593
594                 layer.name = name;
595                 layer.samples = samples;
596
597                 /* If the sample value isn't set yet, check if there is a layer-specific one in the input file. */
598                 if(layer.samples < 1) {
599                         string sample_string = in_spec.get_string_attribute("cycles." + name + ".samples", "");
600                         if(sample_string != "") {
601                                 if(!sscanf(sample_string.c_str(), "%d", &layer.samples)) {
602                                         error = "Failed to parse samples metadata: " + sample_string;
603                                         return false;
604                                 }
605                         }
606                 }
607
608                 if(layer.samples < 1) {
609                         error = string_printf("No sample number specified in the file for layer %s or on the command line", name.c_str());
610                         return false;
611                 }
612
613                 layers.push_back(layer);
614         }
615
616         return true;
617 }
618
619 void DenoiseImage::read_pixels(const DenoiseImageLayer& layer, float *input_pixels)
620 {
621         /* Pixels from center file have already been loaded into pixels.
622          * We copy a subset into the device input buffer with channels reshuffled. */
623         const int *input_to_image_channel = layer.input_to_image_channel.data();
624
625         for(int i = 0; i < width * height; i++) {
626                 for(int j = 0; j < INPUT_NUM_CHANNELS; j++) {
627                         int image_channel = input_to_image_channel[j];
628                         input_pixels[i*INPUT_NUM_CHANNELS + j] = pixels[((size_t)i)*num_channels + image_channel];
629                 }
630         }
631 }
632
633 bool DenoiseImage::read_neighbor_pixels(int neighbor, const DenoiseImageLayer& layer, float *input_pixels)
634 {
635         /* Load pixels from neighboring frames, and copy them into device buffer
636          * with channels reshuffled. */
637         size_t num_pixels = (size_t)width * (size_t)height;
638         array<float> neighbor_pixels(num_pixels * num_channels);
639         if(!in_neighbors[neighbor]->read_image(TypeDesc::FLOAT, neighbor_pixels.data())) {
640                 return false;
641         }
642
643         const int *input_to_image_channel = layer.neighbor_input_to_image_channel[neighbor].data();
644
645         for(int i = 0; i < width * height; i++) {
646                 for(int j = 0; j < INPUT_NUM_CHANNELS; j++) {
647                         int image_channel = input_to_image_channel[j];
648                         input_pixels[i*INPUT_NUM_CHANNELS + j] = neighbor_pixels[((size_t)i)*num_channels + image_channel];
649                 }
650         }
651
652         return true;
653 }
654
655 bool DenoiseImage::load(const string& in_filepath, string& error)
656 {
657         if(!Filesystem::is_regular(in_filepath)) {
658                 error = "Couldn't find file: " + in_filepath;
659                 return false;
660         }
661
662         unique_ptr<ImageInput> in(ImageInput::open(in_filepath));
663         if(!in) {
664                 error = "Couldn't open file: " + in_filepath;
665                 return false;
666         }
667
668         in_spec = in->spec();
669         width = in_spec.width;
670         height = in_spec.height;
671         num_channels = in_spec.nchannels;
672
673         if(!parse_channels(in_spec, error)) {
674                 return false;
675         }
676
677         if(layers.size() == 0) {
678                 error = "Could not find a render layer containing denoising info";
679                 return false;
680         }
681
682         size_t num_pixels = (size_t)width * (size_t)height;
683         pixels.resize(num_pixels * num_channels);
684
685         /* Read all channels into buffer. Reading all channels at once is faster
686          * than individually due to interleaved EXR channel storage. */
687         if(!in->read_image(TypeDesc::FLOAT, pixels.data())) {
688                 error = "Failed to read image: " + in_filepath;
689                 return false;
690         }
691
692         return true;
693 }
694
695 bool DenoiseImage::load_neighbors(const vector<string>& filepaths, const vector<int>& frames, string& error)
696 {
697         if(frames.size() > DENOISE_MAX_FRAMES - 1) {
698                 error = string_printf("Maximum number of neighbors (%d) exceeded\n", DENOISE_MAX_FRAMES - 1);
699                 return false;
700         }
701
702         for(int neighbor = 0; neighbor < frames.size(); neighbor++) {
703                 int frame = frames[neighbor];
704                 const string& filepath = filepaths[frame];
705
706                 if(!Filesystem::is_regular(filepath)) {
707                         error = "Couldn't find neighbor frame: " + filepath;
708                         return false;
709                 }
710
711                 unique_ptr<ImageInput> in_neighbor(ImageInput::open(filepath));
712                 if(!in_neighbor) {
713                         error = "Couldn't open neighbor frame: " + filepath;
714                         return false;
715                 }
716
717                 const ImageSpec &neighbor_spec = in_neighbor->spec();
718                 if(neighbor_spec.width != width || neighbor_spec.height != height) {
719                         error = "Neighbor frame has different dimensions: " + filepath;
720                         return false;
721                 }
722
723                 foreach(DenoiseImageLayer& layer, layers) {
724                         if(!layer.match_channels(neighbor,
725                                                  in_spec.channelnames,
726                                                  neighbor_spec.channelnames))
727                         {
728                                 error = "Neighbor frame misses denoising data passes: " + filepath;
729                                 return false;
730                         }
731                 }
732
733                 in_neighbors.push_back(std::move(in_neighbor));
734         }
735
736         return true;
737 }
738
739 bool DenoiseImage::save_output(const string& out_filepath, string& error)
740 {
741         /* Save image with identical dimensions, channels and metadata. */
742         ImageSpec out_spec = in_spec;
743
744         /* Ensure that the output frame contains sample information even if the input didn't. */
745         for(int i = 0; i < layers.size(); i++) {
746                 string name = "cycles." + layers[i].name + ".samples";
747                 if(!out_spec.find_attribute(name, TypeDesc::STRING)) {
748                         out_spec.attribute(name, TypeDesc::STRING, string_printf("%d", layers[i].samples));
749                 }
750         }
751
752         /* We don't need input anymore at this point, and will possibly
753          * overwrite the same file. */
754         close_input();
755
756         /* Write to temporary file path, so we denoise images in place and don't
757          * risk destroying files when something goes wrong in file saving. */
758         string tmp_filepath = OIIO::Filesystem::temp_directory_path() + "/" + OIIO::Filesystem::unique_path() + ".exr";
759         unique_ptr<ImageOutput> out(ImageOutput::create(tmp_filepath));
760
761         if(!out) {
762                 error = "Failed to open temporary file " + tmp_filepath + " for writing";
763                 return false;
764         }
765
766         /* Open temporary file and write image buffers. */
767         if(!out->open(tmp_filepath, out_spec)) {
768                 error = "Failed to open file " + tmp_filepath + " for writing: " + out->geterror();
769                 return false;
770         }
771
772         bool ok = true;
773         if(!out->write_image(TypeDesc::FLOAT, pixels.data())) {
774                 error = "Failed to write to file " + tmp_filepath + ": " + out->geterror();
775                 ok = false;
776         }
777
778         if(!out->close()) {
779                 error = "Failed to save to file " + tmp_filepath + ": " + out->geterror();
780                 ok = false;
781         }
782
783         out.reset();
784
785         /* Copy temporary file to outputput filepath. */
786         if(ok && !OIIO::Filesystem::rename(tmp_filepath, out_filepath)) {
787                 error = "Failed to save to file " + out_filepath + ": " + out->geterror();
788                 ok = false;
789         }
790
791         if(!ok) {
792                 OIIO::Filesystem::remove(tmp_filepath);
793                 return false;
794         }
795
796         return true;
797 }
798
799 /* File pattern handling and outer loop over frames */
800
801 Denoiser::Denoiser(DeviceInfo& device_info)
802 {
803         samples_override = 0;
804         tile_size = make_int2(64, 64);
805
806         num_frames = 0;
807
808         /* Initialize task scheduler. */
809         TaskScheduler::init();
810
811         /* Initialize device. */
812         DeviceRequestedFeatures req;
813         device = Device::create(device_info, stats, profiler, true);
814         device->load_kernels(req);
815 }
816
817 Denoiser::~Denoiser()
818 {
819         delete device;
820         TaskScheduler::exit();
821 }
822
823 bool Denoiser::run()
824 {
825         assert(input.size() == output.size());
826
827         num_frames = output.size();
828
829         for(int frame = 0; frame < num_frames; frame++) {
830                 /* Skip empty output paths. */
831                 if(output[frame].empty()) {
832                         continue;
833                 }
834
835                 /* Determine neighbor frame numbers that should be used for filtering. */
836                 vector<int> neighbor_frames;
837                 for(int f = frame - params.neighbor_frames; f <= frame + params.neighbor_frames; f++) {
838                         if (f >= 0 && f < num_frames && f != frame) {
839                                 neighbor_frames.push_back(f);
840                         }
841                 }
842
843                 /* Execute task. */
844                 DenoiseTask task(device, this, frame, neighbor_frames);
845                 if(!task.load()) {
846                         error = task.error;
847                         return false;
848                 }
849
850                 if(!task.exec()) {
851                         error = task.error;
852                         return false;
853                 }
854
855                 if(!task.save()) {
856                         error = task.error;
857                         return false;
858                 }
859
860                 task.free();
861         }
862
863         return true;
864 }
865
866 CCL_NAMESPACE_END