| | |
| | |
| | |
| | |
| | |
| | |
| | |
| |
|
| | import os |
| | import re |
| | import numpy as np |
| | import imgui |
| | import PIL.Image |
| | from gui_utils import imgui_utils |
| | from . import renderer |
| | import torch |
| | import torchvision |
| |
|
| | |
| |
|
| | class CaptureWidget: |
| | def __init__(self, viz): |
| | self.viz = viz |
| | self.path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '_screenshots')) |
| | self.dump_image = False |
| | self.dump_gui = False |
| | self.defer_frames = 0 |
| | self.disabled_time = 0 |
| |
|
| | def dump_png(self, image): |
| | viz = self.viz |
| | try: |
| | _height, _width, channels = image.shape |
| | print(viz.result) |
| | assert image.dtype == np.uint8 |
| | os.makedirs(self.path, exist_ok=True) |
| | file_id = 0 |
| | for entry in os.scandir(self.path): |
| | if entry.is_file(): |
| | match = re.fullmatch(r'(\d+).*', entry.name) |
| | if match: |
| | file_id = max(file_id, int(match.group(1)) + 1) |
| | if channels == 1: |
| | pil_image = PIL.Image.fromarray(image[:, :, 0], 'L') |
| | else: |
| | pil_image = PIL.Image.fromarray(image[:, :, :3], 'RGB') |
| | pil_image.save(os.path.join(self.path, f'{file_id:05d}.png')) |
| | np.save(os.path.join(self.path, f'{file_id:05d}.npy'), viz.result.w) |
| | except: |
| | viz.result.error = renderer.CapturedException() |
| |
|
| | @imgui_utils.scoped_by_object_id |
| | def __call__(self, show=True): |
| | viz = self.viz |
| | if show: |
| | with imgui_utils.grayed_out(self.disabled_time != 0): |
| | imgui.text('Capture') |
| | imgui.same_line(viz.label_w) |
| |
|
| | _changed, self.path = imgui_utils.input_text('##path', self.path, 1024, |
| | flags=(imgui.INPUT_TEXT_AUTO_SELECT_ALL | imgui.INPUT_TEXT_ENTER_RETURNS_TRUE), |
| | width=(-1), |
| | help_text='PATH') |
| | if imgui.is_item_hovered() and not imgui.is_item_active() and self.path != '': |
| | imgui.set_tooltip(self.path) |
| | imgui.text(' ') |
| | imgui.same_line(viz.label_w) |
| | if imgui_utils.button('Save image', width=viz.button_w, enabled=(self.disabled_time == 0 and 'image' in viz.result)): |
| | self.dump_image = True |
| | self.defer_frames = 2 |
| | self.disabled_time = 0.5 |
| | imgui.same_line() |
| | if imgui_utils.button('Save GUI', width=viz.button_w, enabled=(self.disabled_time == 0)): |
| | self.dump_gui = True |
| | self.defer_frames = 2 |
| | self.disabled_time = 0.5 |
| |
|
| | self.disabled_time = max(self.disabled_time - viz.frame_delta, 0) |
| | if self.defer_frames > 0: |
| | self.defer_frames -= 1 |
| | elif self.dump_image: |
| | if 'image' in viz.result: |
| | self.dump_png(viz.result.image) |
| | self.dump_image = False |
| | elif self.dump_gui: |
| | viz.capture_next_frame() |
| | self.dump_gui = False |
| | captured_frame = viz.pop_captured_frame() |
| | if captured_frame is not None: |
| | self.dump_png(captured_frame) |
| |
|
| | |
| |
|
