diff --git a/python/taichi/visual/renderer.py b/python/taichi/visual/renderer.py index d91db19f8355b..76a30fa406122 100644 --- a/python/taichi/visual/renderer.py +++ b/python/taichi/visual/renderer.py @@ -1,12 +1,22 @@ from taichi.misc.util import * -import time + +import atexit import os +import time +import numpy + +from concurrent.futures import ThreadPoolExecutor +from subprocess import Popen +from tempfile import mkstemp + +from PIL import Image + import taichi + from taichi.core import tc_core from taichi.misc.util import get_unique_task_id from taichi.visual.post_process import LDRDisplay from taichi.misc.settings import get_num_cores -import cv2 class Renderer(object): @@ -16,6 +26,8 @@ def __init__(self, name=None, output_dir=get_unique_task_id(), overwrite=True, f self.output_dir = taichi.settings.get_output_path(output_dir + '/') self.post_processor = LDRDisplay() self.frame = frame + self.viewer_started = False + self.viewer_process = None try: os.mkdir(self.output_dir) except Exception as e: @@ -54,7 +66,8 @@ def get_full_fn(self, fn): return self.output_dir + fn def write(self, fn): - cv2.imwrite(self.get_full_fn(fn), self.get_output() * 255) + with open(self.get_full_fn(fn), 'wb') as f: + self.get_image_output().save(f) def get_output(self): output = self.c.get_output() @@ -63,9 +76,44 @@ def get_output(self): output = self.post_processor.process(output) return output + def get_image_output(self): + return Image.fromarray((self.get_output() * 255).astype(numpy.uint8), 'RGB') + def show(self): - cv2.imshow('Rendered', self.get_output()) - cv2.waitKey(1) + # allow the user to opt out of the frame viewer by invoking the script + # with --no-viewer in the command line + if '--no-viewer' in sys.argv: + return + + frame_path = self.get_full_fn('current-frame.png') + + # atomic write so watchers don't get a partial image + temp_file, temp_path = mkstemp() + temp_file = os.fdopen(temp_file, 'wb') + self.get_image_output().save(temp_file, format='png') + temp_file.flush() + os.fsync(temp_file.fileno()) + temp_file.close() + os.rename(temp_path, frame_path) + + if not self.viewer_started: + self.viewer_started = True + + pool = ThreadPoolExecutor(max_workers=1) + pool.submit(self.start_viewer, frame_path) + + def end_viewer_process(self): + if self.viewer_process.returncode is not None: + return + + self.viewer_process.terminate() + + def start_viewer(self, frame_path): + path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'viewer.py') + + self.viewer_process = Popen([path, frame_path]) + + atexit.register(self.end_viewer_process) def __getattr__(self, key): return self.c.__getattribute__(key) diff --git a/python/taichi/visual/viewer.py b/python/taichi/visual/viewer.py new file mode 100755 index 0000000000000..a45e2607309da --- /dev/null +++ b/python/taichi/visual/viewer.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python + +import os +import platform +import sys +import Tkinter as tk + +from PIL import Image, ImageTk + +from watchdog.events import FileSystemEventHandler +from watchdog.observers import Observer + + +class EventHandler(FileSystemEventHandler): + + def __init__(self, filename, label, *args, **kwargs): + self.filename = filename + self.label = label + + super(EventHandler, self).__init__(*args, **kwargs) + + def on_created(self, event): + if event.src_path != self.filename: + return + + photo = ImageTk.PhotoImage(Image.open(self.filename)) + + self.label.configure(image=photo) + self.label.image = photo + + +class App(object): + + def __init__(self, filename): + self.filename = filename + + def callback(self): + self.root.quit() + + def run(self): + self.root = tk.Tk() + + self.root.configure(background='black') + self.root.title('Frame Viewer') + self.root.protocol('WM_DELETE_WINDOW', self.callback) + + photo = ImageTk.PhotoImage(Image.open(self.filename)) + + label = tk.Label(self.root, image=photo, background='black') + label.image = photo # keep a reference! + label.pack() + + observer = Observer() + event_handler = EventHandler(self.filename, label) + + observer.schedule(event_handler, os.path.dirname(self.filename)) + observer.start() + + # bring the window to the front when launched + if platform.system != 'Darwin': + self.root.lift() + self.root.attributes('-topmost', True) + self.root.after_idle(self.root.attributes, '-topmost', False) + else: + from Cocoa import NSRunningApplication, NSApplicationActivateIgnoringOtherApps + + app = NSRunningApplication.runningApplicationWithProcessIdentifier(os.getpid()) + app.activateWithOptions(NSApplicationActivateIgnoringOtherApps) + + self.root.mainloop() + + +if __name__ == '__main__': + app = App(sys.argv[1]) + app.run()