From be45d5c9bb497f559402f7a76debcc986caf0fb4 Mon Sep 17 00:00:00 2001
From: evonneng <evonneng@users.noreply.github.com>
Date: Thu, 7 Jul 2022 17:56:22 -0700
Subject: [PATCH] moved logic for interrupting rendering to visutils (#102)

---
 pyrad/graphs/base.py                | 41 +++++++++++++----------------
 pyrad/viewer/server/viewer_utils.py | 16 ++++++++---
 2 files changed, 30 insertions(+), 27 deletions(-)

diff --git a/pyrad/graphs/base.py b/pyrad/graphs/base.py
index 8c4d1b9f4..d65a977ed 100644
--- a/pyrad/graphs/base.py
+++ b/pyrad/graphs/base.py
@@ -18,7 +18,6 @@ The Graph module contains all trainable parameters.
 from abc import abstractmethod
 from collections import defaultdict
 import time
-import traceback
 from typing import Any, Dict, List, Union
 
 import torch
@@ -34,7 +33,6 @@ from pyrad.graphs.modules.ray_generator import RayGenerator
 from pyrad.utils import profiler
 from pyrad.utils.config import GraphConfig
 from pyrad.utils.misc import get_masked_dict, instantiate_from_dict_config, is_not_none
-from pyrad.viewer.server.viewer_utils import CameraChangeException
 
 
 @profiler.time_function
@@ -178,27 +176,24 @@ class Graph(AbstractGraph):
     @torch.no_grad()
     def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: RayBundle):
         """Takes in camera parameters and computes the output of the graph."""
-        try:
-            assert is_not_none(camera_ray_bundle.num_rays_per_chunk)
-            image_height, image_width = camera_ray_bundle.origins.shape[:2]
-            num_rays = len(camera_ray_bundle)
-            outputs = {}
-            outputs_lists = defaultdict(list)
-            for i in range(0, num_rays, camera_ray_bundle.num_rays_per_chunk):
-                start_idx = i
-                end_idx = i + camera_ray_bundle.num_rays_per_chunk
-                ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx)
-                outputs = self.forward_after_ray_generator(ray_bundle)
-                for output_name, output in outputs.items():
-                    outputs_lists[output_name].append(output)
-                time.sleep(0.001)  # visualizer allow thread to switch off
-            for output_name, outputs_list in outputs_lists.items():
-                outputs[output_name] = torch.cat(outputs_list).view(image_height, image_width, -1)
-                time.sleep(0.001)  # visualizer allow thread to switch off
-            self.vis_outputs = outputs
-            return outputs
-        except CameraChangeException:
-            return None
+        assert is_not_none(camera_ray_bundle.num_rays_per_chunk)
+        image_height, image_width = camera_ray_bundle.origins.shape[:2]
+        num_rays = len(camera_ray_bundle)
+        outputs = {}
+        outputs_lists = defaultdict(list)
+        for i in range(0, num_rays, camera_ray_bundle.num_rays_per_chunk):
+            start_idx = i
+            end_idx = i + camera_ray_bundle.num_rays_per_chunk
+            ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx)
+            outputs = self.forward_after_ray_generator(ray_bundle)
+            for output_name, output in outputs.items():
+                outputs_lists[output_name].append(output)
+            time.sleep(0.001)  # visualizer allow thread to switch off
+        for output_name, outputs_list in outputs_lists.items():
+            outputs[output_name] = torch.cat(outputs_list).view(image_height, image_width, -1)
+            time.sleep(0.001)  # visualizer allow thread to switch off
+        self.vis_outputs = outputs
+        return outputs
 
     def get_outputs_for_camera(self, camera: Camera):
         """Get the graph outputs for a Camera."""
diff --git a/pyrad/viewer/server/viewer_utils.py b/pyrad/viewer/server/viewer_utils.py
index 0f72b7fc4..88a73644f 100644
--- a/pyrad/viewer/server/viewer_utils.py
+++ b/pyrad/viewer/server/viewer_utils.py
@@ -85,6 +85,7 @@ class VisualizerState:
             self._render_image_in_viewer(graph)
 
     def _check_interrupt(self, frame, event, arg):
+        """raises interrupt when flag has been set and not already on lowest resolution"""
         if event == "line":
             if self.check_interrupt_vis and self.res_upscale_factor > 1:
                 self.res_upscale_factor = 1
@@ -96,9 +97,8 @@ class VisualizerState:
         if self.vis and step != 0:
             if self.res_upscale_factor == 1:
                 return True
-            else:
-                steps_per_render_image = min(default_steps * self.res_upscale_factor, 100)
-                return step % steps_per_render_image == 0
+            steps_per_render_image = min(default_steps * self.res_upscale_factor, 100)
+            return step % steps_per_render_image == 0
         return False
 
     def _draw_scene_in_viewer(self, image_dataset, dataset_inputs):
@@ -138,11 +138,19 @@ class VisualizerState:
                 self.check_interrupt_vis = True
             time.sleep(0.001)
 
+    def _interruptable_get_outputs_for_camera_ray_bundle(self, graph, camera_ray_bundle: RayBundle):
+        """wrapper around graph function, terminating when interrupting"""
+        try:
+            outputs = graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
+            return outputs
+        except CameraChangeException:
+            return None
+
     @torch.no_grad()
     def _async_get_visualizer_outputs(self, graph, camera_ray_bundle: RayBundle):
         """async getter function for visualizer without returning"""
         with SetTrace(self._check_interrupt):
-            graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
+            self._interruptable_get_outputs_for_camera_ray_bundle(graph, camera_ray_bundle)
         self.check_done_render = True
         self.check_interrupt_vis = False