Skip to content

Commit

Permalink
Improved plotting tools.
Browse files Browse the repository at this point in the history
  • Loading branch information
@frankong
frankong committed Aug 22, 2018
1 parent d5bec09 commit aeb5c8f
Show file tree
Hide file tree
Showing 3 changed files with 143 additions and 74 deletions.
2 changes: 1 addition & 1 deletion sigpy/plot/image.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ def __init__(self, im, x=-1, y=-2, z=None, c=None, hide=False, mode='m', title='
raise TypeError('Image dimension must at least be two, got {im_ndim}'.format(
im_ndim=im.ndim))

self.axim = None
self.im = im
self.fig = plt.figure()
self.ax = self.fig.add_subplot(111)
Expand All @@ -55,7 +56,6 @@ def __init__(self, im, x=-1, y=-2, z=None, c=None, hide=False, mode='m', title='
self.title = title
self.interpolation = interpolation
self.mode = mode
self.axim = None
self.entering_slice = False
self.vmin = None
self.vmax = None
Expand Down
129 changes: 73 additions & 56 deletions sigpy/plot/line.py
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
import os
import uuid
import subprocess
import datetime
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as ani

from tkinter import filedialog
from sigpy.util import prod


Expand All @@ -20,10 +22,15 @@ class Line(object):
<r>: real mode
<i>: imaginary mode
<l>: log mode
<s>: save as png.
<g>: save as gif by traversing current dimension.
<v>: save as mp4 by traversing current dimension.
"""

def __init__(self, arr, x=-1, hide=False, mode='m', title=''):
def __init__(self, arr, x=-1, hide=False, mode='m', title='',
save_basename='Figure', fps=10):
self.arr = arr
self.axarr = None

self.fig = plt.figure()
self.ax = self.fig.add_subplot(111)
Expand All @@ -36,32 +43,19 @@ def __init__(self, arr, x=-1, hide=False, mode='m', title=''):
self.hide = hide
self.title = title
self.mode = mode
self.axarr = None

slices = [0] * (self.ndim - 1) + [slice(None)]
if mode == 'm':
arrv = np.abs(self.arr[slices])
elif mode == 'p':
arrv = np.angle(self.arr[slices])
elif mode == 'r':
arrv = np.real(self.arr[slices])
elif mode == 'i':
arrv = np.imag(self.arr[slices])
elif self.mode == 'l':
eps = 1e-31
arrv = np.log(np.abs(self.arr[slices]) + eps)
self.save_basename = save_basename
self.fps = fps

self.fig.canvas.mpl_disconnect(
self.fig.canvas.manager.key_press_handler_id)
self.fig.canvas.mpl_connect('key_press_event', self.key_press)
self.update_axes()
self.update_line()
self.fig.canvas.draw()
plt.show()

def key_press(self, event):

if event.key == 'up':

if self.d != self.x:
self.slices[self.d] = (
self.slices[self.d] + 1) % self.shape[self.d]
Expand Down Expand Up @@ -123,45 +117,68 @@ def key_press(self, event):
self.update_axes()
self.update_line()
self.fig.canvas.draw()

elif event.key == 's':
file_path = filedialog.asksaveasfilename(filetypes=(("png files", "*.png"),
("pdf files",
"*.pdf"),
("eps files",
"*.eps"),
("svg files",
"*.svg"),
("jpeg files",
"*.jpg"),
("all files", "*.*")))

if not file_path:
return

filename = self.save_basename + \
datetime.datetime.now().strftime(' %Y-%m-%d at %h.%M.%S %p.png')
self.fig.savefig(filename, transparent=True, format='png',
bbox_inches='tight', pad_inches=0)

elif event.key == 'g':
filename = self.save_basename + \
datetime.datetime.now().strftime(' %Y-%m-%d at %h.%M.%S %p.gif')
temp_basename = uuid.uuid4()

bbox = self.fig.get_tightbbox(self.fig.canvas.get_renderer())
for i in range(self.shape[self.d]):
self.slices[self.d] = i

self.update_axes()
self.update_line()
self.fig.canvas.draw()
self.fig.savefig('{} {:05d}.png'.format(temp_basename, i),
format='png', bbox_inches=bbox, pad_inches=0)

subprocess.run(['ffmpeg', '-f', 'image2',
'-s', '{}x{}'.format(int(bbox.width * self.fig.dpi), int(bbox.height * self.fig.dpi)),
'-r', str(self.fps),
'-i', '{} %05d.png'.format(temp_basename),
'-vf', 'palettegen', '{} palette.png'.format(temp_basename)])
subprocess.run(['ffmpeg', '-f', 'image2',
'-s', '{}x{}'.format(int(bbox.width * self.fig.dpi), int(bbox.height * self.fig.dpi)),
'-r', str(self.fps),
'-i', '{} %05d.png'.format(temp_basename),
'-i', '{} palette.png'.format(temp_basename),
'-lavfi', 'paletteuse', filename])

os.remove('{} palette.png'.format(temp_basename))
for i in range(self.shape[self.d]):
os.remove('{} {:05d}.png'.format(temp_basename, i))

elif event.key == 'v':
file_path = filedialog.asksaveasfilename(filetypes=(("mp4 files", "*.mp4"),
("all files", "*.*")))

if not file_path:
return

try:
FFMpegWriter = ani.writers['ffmpeg']
except:
raise ValueError('Does not have FFMPEG installed.')

writer = FFMpegWriter(fps=10)

with writer.saving(self.fig, file_path, 100):
for i in range(self.shape[self.d]):
self.slices[self.d] = i

self.update_axes()
self.update_image()
self.fig.canvas.draw()
writer.grab_frame()

filename = self.save_basename + \
datetime.datetime.now().strftime(' %Y-%m-%d at %h.%M.%S %p.mp4')
temp_basename = uuid.uuid4()

bbox = self.fig.get_tightbbox(self.fig.canvas.get_renderer())
for i in range(self.shape[self.d]):
self.slices[self.d] = i

self.update_axes()
self.update_line()
self.fig.canvas.draw()
self.fig.savefig('{} {:05d}.png'.format(temp_basename, i),
format='png', bbox_inches=bbox, pad_inches=0)

subprocess.run(['ffmpeg', '-f', 'image2',
'-s', '{}x{}'.format(int(bbox.width * self.fig.dpi), int(bbox.height * self.fig.dpi)),
'-r', str(self.fps),
'-i', '{} %05d.png'.format(temp_basename),
'-vf', "scale=trunc(iw/2)*2:trunc(ih/2)*2",
'-vcodec', 'libx264', '-pix_fmt', 'yuv420p', filename])

for i in range(self.shape[self.d]):
os.remove('{} {:05d}.png'.format(temp_basename, i))
else:
return

Expand Down
86 changes: 69 additions & 17 deletions sigpy/plot/scatter.py
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
import os
import uuid
import subprocess
import datetime
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as ani

from tkinter import filedialog
from sigpy.util import prod, move


Expand All @@ -22,7 +24,8 @@ class Scatter(object):
<l>: log mode
"""

def __init__(self, coord, data=None, z=None, hide=False, mode='m', title=''):
def __init__(self, coord, data=None, z=None, hide=False, mode='m', title='',
save_basename='Figure', fps=10):

self.coord = coord
assert coord.shape[-1] == 2
Expand Down Expand Up @@ -52,6 +55,8 @@ def __init__(self, coord, data=None, z=None, hide=False, mode='m', title=''):
self.entering_slice = False
self.vmin = None
self.vmax = None
self.save_basename = save_basename
self.fps = fsp

self.fig.canvas.mpl_disconnect(
self.fig.canvas.manager.key_press_handler_id)
Expand Down Expand Up @@ -125,21 +130,68 @@ def key_press(self, event):
self.update_axes()
self.update_data()
self.fig.canvas.draw()

elif event.key == 's':
file_path = filedialog.asksaveasfilename(filetypes=(("png files", "*.png"),
("pdf files",
"*.pdf"),
("eps files",
"*.eps"),
("svg files",
"*.svg"),
("jpeg files",
"*.jpg"),
("all files", "*.*")))

if not file_path:
return
filename = self.save_basename + \
datetime.datetime.now().strftime(' %Y-%m-%d at %h.%M.%S %p.png')
self.fig.savefig(filename, transparent=True, format='png',
bbox_inches='tight', pad_inches=0)

elif event.key == 'g':
filename = self.save_basename + \
datetime.datetime.now().strftime(' %Y-%m-%d at %h.%M.%S %p.gif')
temp_basename = uuid.uuid4()

bbox = self.fig.get_tightbbox(self.fig.canvas.get_renderer())
for i in range(self.shape[self.d]):
self.slices[self.d] = i

self.update_axes()
self.update_data()
self.fig.canvas.draw()
self.fig.savefig('{} {:05d}.png'.format(temp_basename, i),
format='png', bbox_inches=bbox, pad_inches=0)

subprocess.run(['ffmpeg', '-f', 'image2',
'-s', '{}x{}'.format(int(bbox.width * self.fig.dpi), int(bbox.height * self.fig.dpi)),
'-r', str(self.fps),
'-i', '{} %05d.png'.format(temp_basename),
'-vf', 'palettegen', '{} palette.png'.format(temp_basename)])
subprocess.run(['ffmpeg', '-f', 'image2',
'-s', '{}x{}'.format(int(bbox.width * self.fig.dpi), int(bbox.height * self.fig.dpi)),
'-r', str(self.fps),
'-i', '{} %05d.png'.format(temp_basename),
'-i', '{} palette.png'.format(temp_basename),
'-lavfi', 'paletteuse', filename])

os.remove('{} palette.png'.format(temp_basename))
for i in range(self.shape[self.d]):
os.remove('{} {:05d}.png'.format(temp_basename, i))

elif event.key == 'v':
filename = self.save_basename + \
datetime.datetime.now().strftime(' %Y-%m-%d at %h.%M.%S %p.mp4')
temp_basename = uuid.uuid4()

bbox = self.fig.get_tightbbox(self.fig.canvas.get_renderer())
for i in range(self.shape[self.d]):
self.slices[self.d] = i

self.update_axes()
self.update_data()
self.fig.canvas.draw()
self.fig.savefig('{} {:05d}.png'.format(temp_basename, i),
format='png', bbox_inches=bbox, pad_inches=0)

subprocess.run(['ffmpeg', '-f', 'image2',
'-s', '{}x{}'.format(int(bbox.width * self.fig.dpi), int(bbox.height * self.fig.dpi)),
'-r', str(self.fps),
'-i', '{} %05d.png'.format(temp_basename),
'-vf', "scale=trunc(iw/2)*2:trunc(ih/2)*2",
'-vcodec', 'libx264', '-pix_fmt', 'yuv420p', filename])

for i in range(self.shape[self.d]):
os.remove('{} {:05d}.png'.format(temp_basename, i))

elif (event.key in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'backspace'] and
self.d != self.z):
Expand Down

0 comments on commit aeb5c8f

Please sign in to comment.