diff options
| -rw-r--r-- | chroma/camera.py | 29 |
1 files changed, 26 insertions, 3 deletions
diff --git a/chroma/camera.py b/chroma/camera.py index 9f69f1a..54e8ba7 100644 --- a/chroma/camera.py +++ b/chroma/camera.py @@ -249,7 +249,15 @@ class Camera(multiprocessing.Process): if not os.path.exists(path): break - pygame.image.save(self.screen, path) + try: + pygame.image.save(self.screen, path) + except ImportError: + import Image + mode = 'RGBA' + data = self.screen.get_buffer() + im = Image.frombuffer(mode,self.size,data,'raw',mode,0,1) + im.save(path) + print 'image saved to %s' % path def rotate(self, phi, n): @@ -672,6 +680,8 @@ class EventViewer(Camera): self.rr = RootReader(filename) self.display_mode = EventViewer.CHARGE self.sum_mode = False + self.photon_display_iter = itertools.cycle(['beg','end']) + self.photon_display_mode = self.photon_display_iter.next() def render_particle_track(self): x = 10.0 @@ -680,9 +690,14 @@ class EventViewer(Camera): [0]*3, [0]*3) marker = Solid(pyramid, vacuum, vacuum) + if self.photon_display_mode == 'beg': + photons = self.ev.photons_beg + else: + photons = self.ev.photons_end + geometry = Geometry() - sample_factor = max(1, len(self.ev.photons_beg.pos) / 10000) - for pos in self.ev.photons_beg.pos[::sample_factor]: + sample_factor = max(1, len(photons.pos) / 10000) + for pos in photons.pos[::sample_factor]: geometry.add_solid(marker, displacement=pos, rotation=make_rotation_matrix(np.random.uniform(0,2*np.pi), uniform_sphere())) geometry = create_geometry_from_obj(geometry) @@ -715,6 +730,9 @@ class EventViewer(Camera): def color_hit_pmts(self): self.gpu_geometry.reset_colors() + if self.ev.channels is None: + return + if self.sum_mode: hit = self.sum_hit t = self.sum_t @@ -751,6 +769,11 @@ class EventViewer(Camera): def process_event(self, event): if event.type == KEYDOWN: + if event.key == K_t: + self.photon_display_mode = self.photon_display_iter.next() + self.render_particle_track() + self.update() + if event.key == K_PAGEUP and not self.sum_mode: try: self.ev = self.rr.next() |