class AnimationGui(HasTraits):
figure = Instance(Figure, ())
system = Instance(Catenary, ())
view = View(Group(Group(
Item("object.system.g"),
Item("object.system.length"),
Item("object.system.k"),
Item("object.system.dump"),
),
Item("figure", editor=MPLFigureEditor(toolbar=True)),
show_labels=False,
orientation='horizontal'),
width=800,
height=600,
resizable=True)
def __init__(self, **kw):
super(AnimationGui, self).__init__(**kw)
self.axe = self.figure.add_subplot(111)
def start_timer(self):
self.system.ode_init()
self.line, = self.axe.plot(self.system.x,
-self.system.y,
"-o",
animated=True)
self.axe.set_xlim(-0.1, 1.1)
self.axe.set_ylim(-1, 0.1)
self.ani = FuncAnimation(self.figure,
self.update_figure,
blit=True,
interval=20)
def stop_timer(self):
self.ani._stop()
def update_figure(self, frame):
self.system.ode_step(0.2)
self.line.set_data(self.system.x, -self.system.y)
return self.line,
class PossionDemo(HasTraits):
left_mask = Instance(ImageMaskDrawer)
right_mask = Instance(ImageMaskDrawer)
left_file = File(path.join(FOLDER, "vinci_src.png"))
right_file = File(path.join(FOLDER, "vinci_target.png"))
figure = Instance(Figure, ())
load_button = Button(u"Load")
clear_button = Button(u"Clear")
mix_button = Button(u"Mix")
view = View(Item("left_file"),
Item("right_file"),
VGroup(
HGroup("load_button",
"clear_button",
"mix_button",
show_labels=False),
Group(Item("figure",
editor=MPLFigureEditor(toolbar=False)),
show_labels=False,
orientation='horizontal')),
width=800,
height=600,
resizable=True,
title="Possion Demo")
def __init__(self, **kw):
super(PossionDemo, self).__init__(**kw)
self.left_axe = self.figure.add_subplot(121)
self.right_axe = self.figure.add_subplot(122)
def load_images(self):
self.left_pic = cv2.imread(self.left_file, 1)
self.right_pic = cv2.imread(self.right_file, 1)
shape = [
max(v1, v2) for v1, v2 in zip(self.left_pic.shape[:2],
self.right_pic.shape[:2])
]
self.left_img = self.left_axe.imshow(self.left_pic[::-1, :, ::-1],
origin="lower")
self.left_axe.axis("off")
self.left_mask = ImageMaskDrawer(self.left_axe,
self.left_img,
mask_shape=shape)
self.right_img = self.right_axe.imshow(self.right_pic[::-1, :, ::-1],
origin="lower")
self.right_axe.axis("off")
self.right_mask = ImageMaskDrawer(self.right_axe,
self.right_img,
mask_shape=shape)
self.left_mask.on_trait_change(self.mask_changed, "drawed")
self.right_mask.on_trait_change(self.mask_changed, "drawed")
self.figure.canvas.draw_idle()
def mask_changed(self, obj, name, new):
if obj is self.left_mask:
src, target = self.left_mask, self.right_mask
else:
src, target = self.right_mask, self.left_mask
target.array.fill(0)
target.array[:, :] = src.array[:, :]
target.update()
self.figure.canvas.draw()
def _load_button_fired(self):
self.load_images()
def _mix_button_fired(self):
lh, lw, _ = self.left_pic.shape
rh, rw, _ = self.right_pic.shape
left_mask = self.left_mask.array[-lh:, :lw, -1]
if np.all(left_mask == 0):
return
dx, dy = self.right_mask.get_mask_offset()
result = possion_mix(self.left_pic, self.right_pic, left_mask,
(dx, rh - lh - dy))
self.right_img.set_data(result[::-1, :, ::-1])
self.right_mask.mask_img.set_visible(False)
self.figure.canvas.draw()
def _clear_button_fired(self):
self.left_mask.array.fill(0)
self.right_mask.array.fill(0)
self.left_mask.update()
self.right_mask.update()
self.figure.canvas.draw()
class Demo(HasTraits):
figure = Instance(Figure, ())
timer = Instance(Timer)
hifreq = Range(0, SAMPLING_RATE / 2, SAMPLING_RATE / 2)
spectrum_list = List
buffer_size = Int
traits_view = View(VGroup(HGroup(Item('buffer_size', style="readonly"),
Item('hifreq')),
Item("figure",
editor=MPLFigureEditor(toolbar=True),
show_label=False),
orientation="vertical"),
resizable=True,
title="Audio Spectrum",
width=900,
height=500,
handler=DemoHandler)
def __init__(self, **traits):
super(Demo, self).__init__(**traits)
self._create_plot_component()
self.queue = Queue()
self.finish_event = threading.Event()
self.thread = threading.Thread(target=self.get_audio_data)
self.thread.start()
self.timer = Timer(10, self.on_timer)
self.win_func = np.hanning(NUM_SAMPLES + 1)[:NUM_SAMPLES]
def _hifreq_changed(self):
self.ax1.set_ylim(0, self.hifreq)
def _create_plot_component(self):
self.ax1 = self.figure.add_subplot(311)
self.ax2 = self.figure.add_subplot(312)
self.ax3 = self.figure.add_subplot(313)
self.figure.subplots_adjust(bottom=0.05,
top=0.95,
left=0.05,
right=0.95)
self.spectrogram_data = np.full((NUM_SAMPLES / 2, SPECTROGRAM_LENGTH),
-90)
x2 = float(SPECTROGRAM_LENGTH * NUM_SAMPLES) / float(SAMPLING_RATE)
y2 = SAMPLING_RATE / 2.0
self.image_spectrogram = self.ax1.imshow(self.spectrogram_data,
aspect="auto",
vmin=-90,
vmax=0,
origin="lower",
extent=[0, x2, 0, y2])
self.frequency = np.linspace(0., y2, num=NUM_SAMPLES / 2)
self.time = np.linspace(0.,
float(NUM_SAMPLES) / SAMPLING_RATE,
num=NUM_SAMPLES,
endpoint=False)
for i in range(NUM_LINES):
self.ax2.plot(self.frequency,
np.zeros_like(self.frequency),
color=cm.Blues(float(i + 1) / (NUM_LINES + 1)))
self.line_wave, = self.ax3.plot(self.time, np.zeros_like(self.time))
self.ax3.set_ylim(-1, 1)
self.ax3.set_xlim(0, self.time[-1])
def on_timer(self, *args):
self.buffer_size = self.queue.qsize()
if self.buffer_size > 1:
spectrum, wave = self.queue.get(False)
self.spectrum_list.append(spectrum)
if len(self.spectrum_list) > NUM_LINES: del self.spectrum_list[0]
for i, data in enumerate(self.spectrum_list):
self.ax2.lines[i].set_ydata(data)
self.ax2.relim()
self.ax2.autoscale_view()
self.line_wave.set_ydata(wave)
self.spectrogram_data[:, :-1] = self.spectrogram_data[:, 1:]
self.spectrogram_data[:, -1] = spectrum
self.image_spectrogram.set_data(self.spectrogram_data)
self.figure.canvas.draw()
def get_audio_data(self):
pa = pyaudio.PyAudio()
stream = pa.open(format=pyaudio.paInt16,
channels=1,
rate=SAMPLING_RATE,
input=True,
frames_per_buffer=NUM_SAMPLES)
while not self.finish_event.is_set():
audio_data = np.fromstring(stream.read(NUM_SAMPLES),
dtype=np.short)
normalized_data = (audio_data / 32768.0)
windowed_data = normalized_data * self.win_func
fft_data = np.abs(np.fft.rfft(windowed_data))[:NUM_SAMPLES / 2]
fft_data /= NUM_SAMPLES
np.log10(fft_data, fft_data)
fft_data *= 20
self.queue.put((fft_data, normalized_data))
stream.close()
class FFTDemo(HasTraits):
select = Enum([cls.__name__ for cls in BaseWave.__subclasses__()])
wave = Instance(BaseWave)
# FFT计算所使用的取样点数，这里用一个Enum类型的属性以供用户从列表中选择
fftsize = Enum(256, [(2**x) for x in range(6, 12)])
# 用于显示FFT的结果
peak_list = Str
# 采用多少个频率合成形
N = Range(1, 40, 4)
figure = Instance(Figure, ())
view = View(
HSplit(
VSplit(
VGroup(
Item("select", label="波形类型"),
Item("wave", style="custom", label="波形属性"),
Item("fftsize", label="FFT点数"),
Item("N", label="合成波频率数"),
show_labels=True,
),
Item("peak_list",
style="custom",
show_label=False,
width=100,
height=250)),
VGroup(Item("figure", editor=MPLFigureEditor(toolbar=True)),
show_labels=False),
),
title="FFT演示",
resizable=True,
width=1280,
height=600,
)
def __init__(self, **traits):
super(FFTDemo, self).__init__(**traits)
self.ax1 = self.figure.add_subplot(211)
self.ax2 = self.figure.add_subplot(212)
self.ax1.set_title("FFT")
self.ax2.set_title("Wave")
self.line_peaks, = self.ax1.plot([0], [0], "o")
self.line_wave, = self.ax2.plot([0], [0])
self.line_wave2, = self.ax2.plot([0], [0])
self._select_changed()
def _select_changed(self):
klass = globals()[self.select]
self.wave = klass()
@on_trait_change("wave, wave.[], fftsize")
def draw(self):
fftsize = self.fftsize
x_data = np.arange(0, 1.0, 1.0 / self.fftsize)
func = self.wave.get_func()
# 将func函数的返回值强制转换成float64
y_data = func(x_data).astype(float)
# 计算频谱
fft_parameters = np.fft.fft(y_data) / len(y_data)
self.fft_parameters = fft_parameters
# 计算各个频率的振幅
fft_data = np.clip(
20 * np.log10(np.abs(fft_parameters))[:fftsize / 2 + 1], -120, 120)
self.line_peaks.set_data(np.arange(0, len(fft_data)),
fft_data) # x坐标为频率编号
self.line_wave.set_data(np.arange(0, self.fftsize), y_data)
# 合成波的x坐标为取样点，显示2个周期
self.line_wave2.set_xdata(np.arange(0, 2 * self.fftsize))
# 将振幅大于-80dB的频率输出
peak_index = (fft_data > -80)
peak_value = fft_data[peak_index][:20]
result = []
for f, v in zip(np.flatnonzero(peak_index), peak_value):
result.append("%02d : %g" % (f, v))
self.peak_list = "\n".join(result)
self.ax1.relim()
self.ax1.autoscale_view()
self.plot_sin_combine()
@on_trait_change("N")
def plot_sin_combine(self):
index, data = fft_combine(self.fft_parameters, self.N, 2)
self.line_wave2.set_ydata(data)
self.ax2.relim()
self.ax2.autoscale_view()
if self.figure.canvas is not None:
self.figure.canvas.draw()
class MandelbrotDemo(HasTraits):
color_maps = List
current_map = Str("Blues_r")
loc_info = Str
figure = Instance(Figure)
N = Int(100)
R = Float(10)
cx = Float(0)
cy = Float(0)
d = Float(1.5)
view = View(VGroup(
HGroup(Item("current_map",
label=u"Color map",
editor=EnumEditor(name="object.color_maps")),
"R",
"N",
show_labels=True),
Item("loc_info", show_label=False, style="readonly"),
Group(Item("figure", editor=MPLFigureEditor(toolbar=False)),
show_labels=False,
orientation='horizontal')),
width=SIZE,
height=SIZE + 80,
title=u"Mandelbrot Demo",
resizable=True)
def __init__(self, **kw):
super(MandelbrotDemo, self).__init__(**kw)
self.array = np.zeros((SIZE, SIZE))
self.img = self.figure.figimage(self.array, cmap=self.current_map)
def _figure_default(self):
return Figure(figsize=(SIZE / 100.0, SIZE / 100.0), dpi=100)
def _color_maps_default(self):
return sorted(cm.cmap_d.keys())
def init_plot(self):
self.figure.canvas.mpl_connect('button_press_event', self.on_press)
self.figure.canvas.mpl_connect('motion_notify_event', self.on_move)
self.figure.canvas.mpl_connect('scroll_event', self.on_scroll)
self.update_plot()
def _current_map_changed(self):
self.img.set_cmap(self.current_map)
self.update_plot()
@on_trait_change("R, N")
def update_plot(self):
mandelbrot(self.cx,
self.cy,
self.d,
out=self.array,
n=self.N,
R=self.R)
self.loc_info = "%g, %g, d=%g" % (self.cx, self.cy, self.d)
self.img.set_data(self.array)
self.figure.canvas.draw()
def on_press(self, evt):
if evt.button == 1:
self.last_pos = evt.x, evt.y, self.cx, self.cy
def on_move(self, evt):
if evt.button != 1:
return
x, y, cx, cy = self.last_pos
dx = (x - evt.x) * 2 * self.d / SIZE
dy = (evt.y - y) * 2 * self.d / SIZE
self.cx = cx + dx
self.cy = cy + dy
self.update_plot()
def on_scroll(self, evt):
x0, x1, y0, y1 = self.cx - self.d, self.cx + self.d, self.cy - self.d, self.cy + self.d
x = x0 + float(evt.x) / SIZE * 2 * self.d
y = y0 + float(SIZE - evt.y) / SIZE * 2 * self.d
if evt.step < 0:
d2 = self.d * 1.2
else:
d2 = self.d / 1.2
scale = d2 / self.d
self.cx = (2 * x + (x0 + x1 - 2 * x) * scale) / 2
self.cy = (2 * y + (y0 + y1 - 2 * y) * scale) / 2
self.d = d2
self.update_plot()