def localize():
    # Global variables that may be set in this function
    global switch_beamforming
    global DO_BEAMFORM
    global done
    # Setup search space
    source_plane = OrientedSourcePlane(SOURCE_PLANE_NORMAL, 
                                       SOURCE_PLANE_UP,
                                       SOURCE_PLANE_OFFSET)
    space = SearchSpace(MIC_LOC, CAMERA_LOC, [source_plane], MIC_FORWARD, MIC_ABOVE)
                                       
    # Setup pyaudio instances
    pa = pyaudio.PyAudio()
    helper = AudioHelper(pa)
    listener = CommandListener()
    plot_manager = PlotManager()
    #localizer = GridTrackingLocalizer(mic_positions=mic_layout,
    #                                  search_space=space,
    #                                  source_cov=SOURCE_LOCATION_COV,
    #                                  dft_len=FFT_LENGTH,
    #                                  sample_rate=SAMPLE_RATE,
    #                                  n_theta=N_THETA,
    #                                  n_phi=N_PHI)
    localizer = KalmanTrackingLocalizer(mic_positions=mic_layout,
                                      search_space=space,
                                      mic_forward=MIC_FORWARD,
                                      mic_above=MIC_ABOVE,
                                      trans_mat=STATE_TRANSITION_MAT,
                                      state_cov=STATE_TRANSITION_MAT,
                                      emission_mat=EMISSION_MAT,
                                      emission_cov=EMISSION_COV,
                                      dft_len=FFT_LENGTH,
                                      sample_rate=SAMPLE_RATE,
                                      n_theta=N_THETA,
                                      n_phi=N_PHI)
    beamformer = BeamFormer(mic_layout, SAMPLE_RATE)

    # Setup STFT object
    stft = StftManager(dft_length=FFT_LENGTH,
                       window_length=WINDOW_LENGTH,
                       hop_length=HOP_LENGTH,
                       use_window_fcn=True,
                       n_channels=NUM_CHANNELS_IN,
                       dtype=DATA_TYPE)

    # Setup devices
    in_device = helper.get_input_device_from_user()
    if PLAY_AUDIO:
        out_device = helper.get_output_device_from_user()
    else:
        out_device = helper.get_default_output_device_info()

    # Setup streams
    in_stream = pa.open(rate=SAMPLE_RATE,
                        channels=NUM_CHANNELS_IN,
                        format=SAMPLE_TYPE,
                        frames_per_buffer=FRAMES_PER_BUF,
                        input=True,
                        input_device_index=int(in_device['index']),
                        stream_callback=read_in_data)
    out_stream = pa.open(rate=SAMPLE_RATE,
                         channels=NUM_CHANNELS_OUT,
                         format=SAMPLE_TYPE,
                         output=True,
                         frames_per_buffer=FRAMES_PER_BUF,
                         output_device_index=int(out_device['index']),
                         stream_callback=write_out_data)

    # Start recording/playing back
    in_stream.start_stream()
    out_stream.start_stream()

    # Start thread to check for user quit
    listener.start_polling()

    # Setup directions and alignment matrices
    direcs = localizer.get_directions()
    align_mats = localizer.get_pos_align_mat()

    # Plotting setup
    if PLOT_POLAR:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='polar')
        ax.set_rlim(0, 1)
        plt.show(block=False)
        # Setup space for plotting in new coordinates
        spher_coords = localizer.get_spher_directions()
        theta = spher_coords[1, :]
        pol_plot, = plt.plot(theta, np.zeros(theta.shape))
        post_plot, = plt. plot(theta, np.zeros(theta.shape), 'green')
        ax.set_ylim(0, 1)
        if DO_BEAMFORM:
            pol_beam_plot, = plt.plot(theta, np.ones(theta.shape), 'red')
    if PLOT_CARTES:
        fig = plt.figure()
        ax = plotting.get_halfpage_axis(fig)
        #ax = fig.add_subplot(111)
        plt.show(block=False)
        # Setup space for plotting in new coordinates
        spher_coords = localizer.get_spher_directions()
        theta = spher_coords[1, :]
        theta = np.linspace(0, 1, theta.shape[0])
        gcc_plots = []
        gcc_shaping_vals = [1, 2, 3, 4, 5]
        for i in gcc_shaping_vals:
            plot, = plt.plot(theta, np.zeros(theta.shape))
            gcc_plots.append(plot)
        pol_plot, = plt.plot(theta, np.zeros(theta.shape), 'r--')
        post_plot, = plt. plot(theta, np.zeros(theta.shape), 'b')
        ax.set_ylim(0, 1.2)
        ax.set_xlim(0, 1)  # Normalized
        #ax.set_xlabel('Angle $\left(\\frac{1}{\pi}\\right)$')
        #ax.set_ylabel('Normalized GCC')
        if DO_BEAMFORM:
            pol_beam_plot, = plt.plot(theta, np.ones(theta.shape), 'red')
    if PLOT_2D:
        n_past_samples = 200
        filter_plot = FilterPlot(N_THETA, n_past_samples, 2)
        save_plot = filter_plot # For saving figures
    if VIDEO_OVERLAY:
        vc = cv2.VideoCapture(0)
        video_handle, video_plot = setup_video_handle(720, 1280)
        plt.show(block=False)
    if EXTERNAL_PLOT:
        fig = plt.figure()
        ax = fig.add_subplot(111)
        plt.show(block=False)

    count = 0
    try:
        while in_stream.is_active() or out_stream.is_active():
            done = listener.quit()
            data_available = in_buf.wait_for_read(WINDOW_LENGTH, TIMEOUT)
            if data_available:
                if listener.switch_beamforming():
                    DO_BEAMFORM = not DO_BEAMFORM
                if listener.savefig():
                    plot_manager.savefig(save_plot.get_figure())
                # Get data from circular buffer
                data = in_buf.read_samples(WINDOW_LENGTH)
                # Perform an stft
                stft.performStft(data)
                # Process dfts from windowed segments of input
                dfts = stft.getDFTs()
                rffts = mat.to_all_real_matlab_format(dfts)
                gccs = []
                #for k in gcc_shaping_vals:
                #    d, energy = localizer.get_distribution_real(
                #            rffts[:, :, 0], 'mcc', k) # Use first hop
                #    gccs.append(d)
                d, energy = localizer.get_distribution_real(rffts[:, :, 0], 'beam') # Use first hop
                def w(cpmat):
                    cpmat /= (np.abs(cpmat + consts.EPS))
                    return cpmat
                post = localizer.get_distribution(rffts[:, :, 0], 'beam')
                #post, bla = localizer.get_distribution_real(rffts[:, :, 0], 'mcc')

                #post = localizer.get_distribution(rffts[:, :, 0])
                ind = np.argmax(d)
                u = 1.5 * direcs[:, ind]  # Direction of arrival
                #if energy < 500:
                    #continue

                # Do beam forming
                if DO_BEAMFORM:
                    align_mat = align_mats[:, :, ind]
                    filtered = beamformer.filter_real(rffts, align_mat)
                    mat.set_dfts_real(dfts, filtered, n_channels=2)

                # Take care of plotting
                if count % 1 == 0:
                    if PLOT_POLAR or PLOT_CARTES:
                        dist = d
                        #dist -= np.min(dist)
                        dist = localizer.to_spher_grid(dist)
                        print post.shape
                        post = localizer.to_spher_grid(post) * 50
                        dist /= np.max(dist)
                        if np.max(dist) > 1:
                          dist /= np.max(dist)
                        if np.max(post) > 1:
                          post /= np.max(post)
                        pol_plot.set_ydata(dist[0, :])
                        post_plot.set_ydata(post[0, :])
                        #for i, plot in enumerate(gcc_plots):
                        #    gcc = gccs[i]
                        #    gcc /= (np.max(gcc) + consts.EPS)
                        #    plot.set_ydata(gccs[i])
                        if DO_BEAMFORM:
                            # Get beam plot
                            freq = 2500.  # Hz
                            response = beamformer.get_beam(
                                align_mat, align_mats, rffts, freq
                            )
                            response = localizer.to_spher_grid(response)
                            if np.max(response) > 1:
                                response /= np.max(response)
                            pol_beam_plot.set_ydata(response[-1, :])
                        plt.draw()
                    if PLOT_2D:
                        dist = localizer.to_spher_grid(d)
                        p = localizer.to_spher_grid(post)
                        est1 = THETA_SPACE[np.argmax(p)]
                        est2 = THETA_SPACE[np.argmax(dist)]
                        filter_plot.update(dist, [est1, est2])
                    if VIDEO_OVERLAY:
                        post /= np.max(post + consts.EPS)
                        dist = d - np.min(d)
                        dist = dist / np.max(dist + consts.EPS)
                        _, cvimage = vc.read()
                        overlay_distribution(video_handle, video_plot, cvimage, dist[::-1])
                        plt.draw()
                count += 1

                # Get the istft of the processed data
                if PLAY_AUDIO or RECORD_AUDIO:
                    new_data = stft.performIStft()
                    new_data = out_buf.reduce_channels(new_data, NUM_CHANNELS_IN, NUM_CHANNELS_OUT)
                    # Write out the new, altered data
                    if PLAY_AUDIO:
                        if out_buf.get_available_write() >= WINDOW_LENGTH:
                            out_buf.write_samples(new_data)
                    if RECORD_AUDIO:
                        if record_buf.get_available_write() >= WINDOW_LENGTH:
                            record_buf.write_samples(new_data)


    except KeyboardInterrupt:
        print "Program interrupted"
        listener.set_quit(True)


    print "Cleaning up"
    in_stream.stop_stream()
    in_stream.close()
    out_stream.stop_stream()
    out_stream.close()
    pa.terminate()

    # Take care of output file
    if RECORD_AUDIO:
        print "Writing output file"
        make_wav()

    print "Done"
def localize():
    global switch_beamforming
    global DO_BEAMFORM
    # Setup search space
    source_plane = OrientedSourcePlane(SOURCE_PLANE_NORMAL, 
                                       SOURCE_PLANE_UP,
                                       SOURCE_PLANE_OFFSET)
    space = SearchSpace(MIC_LOC, CAMERA_LOC, [source_plane])
                                       
    # Setup pyaudio instances
    pa = pyaudio.PyAudio()
    helper = AudioHelper(pa)
    localizer = KalmanTrackingLocalizer(mic_positions=mic_layout,
                                      search_space=space,
                                      mic_forward=MIC_FORWARD,
                                      mic_above=MIC_ABOVE,
                                      trans_mat=STATE_TRANSITION_MAT,
                                      state_cov=STATE_TRANSITION_MAT,
                                      emission_mat=EMISSION_MAT,
                                      emission_cov=EMISSION_COV,
                                      dft_len=FFT_LENGTH,
                                      sample_rate=SAMPLE_RATE,
                                      n_theta=N_THETA,
                                      n_phi=N_PHI)
    beamformer = BeamFormer(mic_layout, SAMPLE_RATE)

    # Setup STFT object
    stft = StftManager(dft_length=FFT_LENGTH,
                       window_length=WINDOW_LENGTH,
                       hop_length=HOP_LENGTH,
                       use_window_fcn=True,
                       n_channels=NUM_CHANNELS_IN,
                       dtype=DATA_TYPE)

    # Setup devices
    in_device = helper.get_input_device_from_user()
    if PLAY_AUDIO:
        out_device = helper.get_output_device_from_user()
    else:
        out_device = helper.get_default_output_device_info()

    # Setup streams
    in_stream = pa.open(rate=SAMPLE_RATE,
                        channels=NUM_CHANNELS_IN,
                        format=SAMPLE_TYPE,
                        frames_per_buffer=FRAMES_PER_BUF,
                        input=True,
                        input_device_index=int(in_device['index']),
                        stream_callback=read_in_data)
    out_stream = pa.open(rate=SAMPLE_RATE,
                         channels=NUM_CHANNELS_OUT,
                         format=SAMPLE_TYPE,
                         output=True,
                         frames_per_buffer=FRAMES_PER_BUF,
                         output_device_index=int(out_device['index']),
                         stream_callback=write_out_data)

    # Start recording/playing back
    in_stream.start_stream()
    out_stream.start_stream()

    # Start thread to check for user quit
    quit_thread = threading.Thread(target=check_for_quit)
    quit_thread.start()

    # Setup directions and alignment matrices
    direcs = localizer.get_directions()
    align_mats = localizer.get_pos_align_mat()

    # Plotting setup
    if PLOT_POLAR:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='polar')
        ax.set_rlim(0, 1)
        plt.show(block=False)
        # Setup space for plotting in new coordinates
        spher_coords = localizer.get_spher_directions()
        theta = spher_coords[1, :]
        pol_plot, = plt.plot(theta, np.ones(theta.shape))
        post_plot, = plt. plot(theta, np.ones(theta.shape), 'green')
        ax.set_ylim(0, 1)
        if DO_BEAMFORM:
            pol_beam_plot, = plt.plot(theta, np.ones(theta.shape), 'red')
    if PLOT_CARTES:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        plt.show(block=False)
        x = localizer.to_spher_grid(direcs[0, :])
        y = localizer.to_spher_grid(direcs[1, :])
        z = localizer.to_spher_grid(direcs[2, :])
        #scat = ax.scatter(x, y, z, s=100)
    if EXTERNAL_PLOT:
        fig = plt.figure()
        ax = fig.add_subplot(111)
        plt.show(block=False)

    count = 0
    try:
        global done
        while in_stream.is_active() or out_stream.is_active():
            data_available = in_buf.wait_for_read(WINDOW_LENGTH, TIMEOUT)
            if data_available:
                if switch_beamforming:
                    DO_BEAMFORM = not DO_BEAMFORM
                    switch_beamforming = False
                    # Get data from the circular buffer
                data = in_buf.read_samples(WINDOW_LENGTH)
                # Perform an stft
                stft.performStft(data)
                # Process dfts from windowed segments of input
                dfts = stft.getDFTs()
                rffts = mat.to_all_real_matlab_format(dfts)
                d, energy = localizer.get_distribution_real(rffts[:, :, 0], 'gcc') # Use first hop
                post = localizer.get_distribution(rffts[:, :, 0])
                ind = np.argmax(post)
                u = 1.5 * direcs[:, ind]  # Direction of arrival
                #if energy < 500:
                #    continue

                # Do beam forming
                if DO_BEAMFORM:
                    align_mat = align_mats[:, :, ind]
                    filtered = beamformer.filter_real(rffts, align_mat)
                    mat.set_dfts_real(dfts, filtered, n_channels=2)

                # Take care of plotting
                if count % 1 == 0:
                    if PLOT_CARTES:
                        ax.cla()
                        ax.grid(False)
                        d = localizer.to_spher_grid(post / (np.max(post) + consts.EPS))
                        #d = localizer.to_spher_grid(d / (np.max(d) + consts.EPS))
                        ax.scatter(x, y, z, c=d, s=40)
                        #ax.plot_surface(x, y, z, rstride=1, cstride=1, facecolor=plt.cm.gist_heat(d))
                        ax.plot([0, u[0]], [0, u[1]], [0, u[2]], c='black', linewidth=3)
                        if DO_BEAMFORM:
                            if np.max(np.abs(response)) > 1:
                                response /= np.max(np.abs(response))
                            X = response * x
                            Y = response * y
                            Z = response * z
                            ax.plot_surface(X, Y, Z, rstride=1, cstride=1, color='white')
                        ax.set_xlim(-1, 1)
                        ax.set_ylim(-1, 1)
                        ax.set_zlim(0, 1)
                        #ax.view_init(90, -90)
                        fig.canvas.draw()
                    if PLOT_2D:
                        # Get unconditional distribution
                        dist = localizer.to_spher_grid(d)
                        dist -= np.min(dist)
                        dist /= (np.sum(dist) + consts.EPS)
                        sample_mat[:, :-1] = sample_mat[:, 1:]
                        sample_mat[:, -1] = dist
                        # Get kalman estimate
                        maxind = np.argmax(post)
                        estimate_mat[:-1] = estimate_mat[1:]
                        estimate_mat[-1] = maxind
                        plot_2d.set_array(sample_mat)
                        state_est_plot.set_ydata(estimate_mat)
                        plt.draw()
                count += 1

                # Get the istft of the processed data
                if PLAY_AUDIO or RECORD_AUDIO:
                    new_data = stft.performIStft()
                    new_data = out_buf.reduce_channels(new_data, NUM_CHANNELS_IN, NUM_CHANNELS_OUT)
                    # Write out the new, altered data
                    if PLAY_AUDIO:
                        if out_buf.get_available_write() >= WINDOW_LENGTH:
                            out_buf.write_samples(new_data)
                    if RECORD_AUDIO:
                        if record_buf.get_available_write() >= WINDOW_LENGTH:
                            record_buf.write_samples(new_data)


    except KeyboardInterrupt:
        print "Program interrupted"
        done = True


    print "Cleaning up"
    in_stream.stop_stream()
    in_stream.close()
    out_stream.stop_stream()
    out_stream.close()
    pa.terminate()

    # Take care of output file
    if RECORD_AUDIO:
        print "Writing output file"
        make_wav()

    print "Done"
def localize():
    # Global variables that may be set in this function
    global switch_beamforming
    global DO_BEAMFORM
    global done
    # Setup search space
    source_plane = OrientedSourcePlane(SOURCE_PLANE_NORMAL, SOURCE_PLANE_UP,
                                       SOURCE_PLANE_OFFSET)
    space = SearchSpace(MIC_LOC, CAMERA_LOC, [source_plane], MIC_FORWARD,
                        MIC_ABOVE)

    # Setup pyaudio instances
    pa = pyaudio.PyAudio()
    helper = AudioHelper(pa)
    listener = CommandListener()
    plot_manager = PlotManager()
    #localizer = GridTrackingLocalizer(mic_positions=mic_layout,
    #                                  search_space=space,
    #                                  source_cov=SOURCE_LOCATION_COV,
    #                                  dft_len=FFT_LENGTH,
    #                                  sample_rate=SAMPLE_RATE,
    #                                  n_theta=N_THETA,
    #                                  n_phi=N_PHI)
    localizer = KalmanTrackingLocalizer(mic_positions=mic_layout,
                                        search_space=space,
                                        mic_forward=MIC_FORWARD,
                                        mic_above=MIC_ABOVE,
                                        trans_mat=STATE_TRANSITION_MAT,
                                        state_cov=STATE_TRANSITION_MAT,
                                        emission_mat=EMISSION_MAT,
                                        emission_cov=EMISSION_COV,
                                        dft_len=FFT_LENGTH,
                                        sample_rate=SAMPLE_RATE,
                                        n_theta=N_THETA,
                                        n_phi=N_PHI)
    beamformer = BeamFormer(mic_layout, SAMPLE_RATE)

    # Setup STFT object
    stft = StftManager(dft_length=FFT_LENGTH,
                       window_length=WINDOW_LENGTH,
                       hop_length=HOP_LENGTH,
                       use_window_fcn=True,
                       n_channels=NUM_CHANNELS_IN,
                       dtype=DATA_TYPE)

    # Setup devices
    in_device = helper.get_input_device_from_user()
    if PLAY_AUDIO:
        out_device = helper.get_output_device_from_user()
    else:
        out_device = helper.get_default_output_device_info()

    # Setup streams
    in_stream = pa.open(rate=SAMPLE_RATE,
                        channels=NUM_CHANNELS_IN,
                        format=SAMPLE_TYPE,
                        frames_per_buffer=FRAMES_PER_BUF,
                        input=True,
                        input_device_index=int(in_device['index']),
                        stream_callback=read_in_data)
    out_stream = pa.open(rate=SAMPLE_RATE,
                         channels=NUM_CHANNELS_OUT,
                         format=SAMPLE_TYPE,
                         output=True,
                         frames_per_buffer=FRAMES_PER_BUF,
                         output_device_index=int(out_device['index']),
                         stream_callback=write_out_data)

    # Start recording/playing back
    in_stream.start_stream()
    out_stream.start_stream()

    # Start thread to check for user quit
    listener.start_polling()

    # Setup directions and alignment matrices
    direcs = localizer.get_directions()
    align_mats = localizer.get_pos_align_mat()

    # Plotting setup
    if PLOT_POLAR:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='polar')
        ax.set_rlim(0, 1)
        plt.show(block=False)
        # Setup space for plotting in new coordinates
        spher_coords = localizer.get_spher_directions()
        theta = spher_coords[1, :]
        pol_plot, = plt.plot(theta, np.zeros(theta.shape))
        post_plot, = plt.plot(theta, np.zeros(theta.shape), 'green')
        ax.set_ylim(0, 1)
        if DO_BEAMFORM:
            pol_beam_plot, = plt.plot(theta, np.ones(theta.shape), 'red')
    if PLOT_CARTES:
        fig = plt.figure()
        ax = plotting.get_halfpage_axis(fig)
        #ax = fig.add_subplot(111)
        plt.show(block=False)
        # Setup space for plotting in new coordinates
        spher_coords = localizer.get_spher_directions()
        theta = spher_coords[1, :]
        theta = np.linspace(0, 1, theta.shape[0])
        gcc_plots = []
        gcc_shaping_vals = [1, 2, 3, 4, 5]
        for i in gcc_shaping_vals:
            plot, = plt.plot(theta, np.zeros(theta.shape))
            gcc_plots.append(plot)
        pol_plot, = plt.plot(theta, np.zeros(theta.shape), 'r--')
        post_plot, = plt.plot(theta, np.zeros(theta.shape), 'b')
        ax.set_ylim(0, 1.2)
        ax.set_xlim(0, 1)  # Normalized
        #ax.set_xlabel('Angle $\left(\\frac{1}{\pi}\\right)$')
        #ax.set_ylabel('Normalized GCC')
        if DO_BEAMFORM:
            pol_beam_plot, = plt.plot(theta, np.ones(theta.shape), 'red')
    if PLOT_2D:
        n_past_samples = 200
        filter_plot = FilterPlot(N_THETA, n_past_samples, 2)
        save_plot = filter_plot  # For saving figures
    if VIDEO_OVERLAY:
        vc = cv2.VideoCapture(0)
        video_handle, video_plot = setup_video_handle(720, 1280)
        plt.show(block=False)
    if EXTERNAL_PLOT:
        fig = plt.figure()
        ax = fig.add_subplot(111)
        plt.show(block=False)

    count = 0
    try:
        while in_stream.is_active() or out_stream.is_active():
            done = listener.quit()
            data_available = in_buf.wait_for_read(WINDOW_LENGTH, TIMEOUT)
            if data_available:
                if listener.switch_beamforming():
                    DO_BEAMFORM = not DO_BEAMFORM
                if listener.savefig():
                    plot_manager.savefig(save_plot.get_figure())
                # Get data from circular buffer
                data = in_buf.read_samples(WINDOW_LENGTH)
                # Perform an stft
                stft.performStft(data)
                # Process dfts from windowed segments of input
                dfts = stft.getDFTs()
                rffts = mat.to_all_real_matlab_format(dfts)
                gccs = []
                #for k in gcc_shaping_vals:
                #    d, energy = localizer.get_distribution_real(
                #            rffts[:, :, 0], 'mcc', k) # Use first hop
                #    gccs.append(d)
                d, energy = localizer.get_distribution_real(
                    rffts[:, :, 0], 'beam')  # Use first hop

                def w(cpmat):
                    cpmat /= (np.abs(cpmat + consts.EPS))
                    return cpmat

                post = localizer.get_distribution(rffts[:, :, 0], 'beam')
                #post, bla = localizer.get_distribution_real(rffts[:, :, 0], 'mcc')

                #post = localizer.get_distribution(rffts[:, :, 0])
                ind = np.argmax(d)
                u = 1.5 * direcs[:, ind]  # Direction of arrival
                #if energy < 500:
                #continue

                # Do beam forming
                if DO_BEAMFORM:
                    align_mat = align_mats[:, :, ind]
                    filtered = beamformer.filter_real(rffts, align_mat)
                    mat.set_dfts_real(dfts, filtered, n_channels=2)

                # Take care of plotting
                if count % 1 == 0:
                    if PLOT_POLAR or PLOT_CARTES:
                        dist = d
                        #dist -= np.min(dist)
                        dist = localizer.to_spher_grid(dist)
                        print post.shape
                        post = localizer.to_spher_grid(post) * 50
                        dist /= np.max(dist)
                        if np.max(dist) > 1:
                            dist /= np.max(dist)
                        if np.max(post) > 1:
                            post /= np.max(post)
                        pol_plot.set_ydata(dist[0, :])
                        post_plot.set_ydata(post[0, :])
                        #for i, plot in enumerate(gcc_plots):
                        #    gcc = gccs[i]
                        #    gcc /= (np.max(gcc) + consts.EPS)
                        #    plot.set_ydata(gccs[i])
                        if DO_BEAMFORM:
                            # Get beam plot
                            freq = 2500.  # Hz
                            response = beamformer.get_beam(
                                align_mat, align_mats, rffts, freq)
                            response = localizer.to_spher_grid(response)
                            if np.max(response) > 1:
                                response /= np.max(response)
                            pol_beam_plot.set_ydata(response[-1, :])
                        plt.draw()
                    if PLOT_2D:
                        dist = localizer.to_spher_grid(d)
                        p = localizer.to_spher_grid(post)
                        est1 = THETA_SPACE[np.argmax(p)]
                        est2 = THETA_SPACE[np.argmax(dist)]
                        filter_plot.update(dist, [est1, est2])
                    if VIDEO_OVERLAY:
                        post /= np.max(post + consts.EPS)
                        dist = d - np.min(d)
                        dist = dist / np.max(dist + consts.EPS)
                        _, cvimage = vc.read()
                        overlay_distribution(video_handle, video_plot, cvimage,
                                             dist[::-1])
                        plt.draw()
                count += 1

                # Get the istft of the processed data
                if PLAY_AUDIO or RECORD_AUDIO:
                    new_data = stft.performIStft()
                    new_data = out_buf.reduce_channels(new_data,
                                                       NUM_CHANNELS_IN,
                                                       NUM_CHANNELS_OUT)
                    # Write out the new, altered data
                    if PLAY_AUDIO:
                        if out_buf.get_available_write() >= WINDOW_LENGTH:
                            out_buf.write_samples(new_data)
                    if RECORD_AUDIO:
                        if record_buf.get_available_write() >= WINDOW_LENGTH:
                            record_buf.write_samples(new_data)

    except KeyboardInterrupt:
        print "Program interrupted"
        listener.set_quit(True)

    print "Cleaning up"
    in_stream.stop_stream()
    in_stream.close()
    out_stream.stop_stream()
    out_stream.close()
    pa.terminate()

    # Take care of output file
    if RECORD_AUDIO:
        print "Writing output file"
        make_wav()

    print "Done"
Example #4
0
def localize():
    global switch_beamforming
    global DO_BEAMFORM
    # Setup search space
    source_plane = OrientedSourcePlane(SOURCE_PLANE_NORMAL, SOURCE_PLANE_UP,
                                       SOURCE_PLANE_OFFSET)
    space = SearchSpace(MIC_LOC, CAMERA_LOC, [source_plane])

    # Setup pyaudio instances
    pa = pyaudio.PyAudio()
    helper = AudioHelper(pa)
    localizer = KalmanTrackingLocalizer(mic_positions=mic_layout,
                                        search_space=space,
                                        mic_forward=MIC_FORWARD,
                                        mic_above=MIC_ABOVE,
                                        trans_mat=STATE_TRANSITION_MAT,
                                        state_cov=STATE_TRANSITION_MAT,
                                        emission_mat=EMISSION_MAT,
                                        emission_cov=EMISSION_COV,
                                        dft_len=FFT_LENGTH,
                                        sample_rate=SAMPLE_RATE,
                                        n_theta=N_THETA,
                                        n_phi=N_PHI)
    beamformer = BeamFormer(mic_layout, SAMPLE_RATE)

    # Setup STFT object
    stft = StftManager(dft_length=FFT_LENGTH,
                       window_length=WINDOW_LENGTH,
                       hop_length=HOP_LENGTH,
                       use_window_fcn=True,
                       n_channels=NUM_CHANNELS_IN,
                       dtype=DATA_TYPE)

    # Setup devices
    in_device = helper.get_input_device_from_user()
    if PLAY_AUDIO:
        out_device = helper.get_output_device_from_user()
    else:
        out_device = helper.get_default_output_device_info()

    # Setup streams
    in_stream = pa.open(rate=SAMPLE_RATE,
                        channels=NUM_CHANNELS_IN,
                        format=SAMPLE_TYPE,
                        frames_per_buffer=FRAMES_PER_BUF,
                        input=True,
                        input_device_index=int(in_device['index']),
                        stream_callback=read_in_data)
    out_stream = pa.open(rate=SAMPLE_RATE,
                         channels=NUM_CHANNELS_OUT,
                         format=SAMPLE_TYPE,
                         output=True,
                         frames_per_buffer=FRAMES_PER_BUF,
                         output_device_index=int(out_device['index']),
                         stream_callback=write_out_data)

    # Start recording/playing back
    in_stream.start_stream()
    out_stream.start_stream()

    # Start thread to check for user quit
    quit_thread = threading.Thread(target=check_for_quit)
    quit_thread.start()

    # Setup directions and alignment matrices
    direcs = localizer.get_directions()
    align_mats = localizer.get_pos_align_mat()

    # Plotting setup
    if PLOT_POLAR:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='polar')
        ax.set_rlim(0, 1)
        plt.show(block=False)
        # Setup space for plotting in new coordinates
        spher_coords = localizer.get_spher_directions()
        theta = spher_coords[1, :]
        pol_plot, = plt.plot(theta, np.ones(theta.shape))
        post_plot, = plt.plot(theta, np.ones(theta.shape), 'green')
        ax.set_ylim(0, 1)
        if DO_BEAMFORM:
            pol_beam_plot, = plt.plot(theta, np.ones(theta.shape), 'red')
    if PLOT_CARTES:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        plt.show(block=False)
        x = localizer.to_spher_grid(direcs[0, :])
        y = localizer.to_spher_grid(direcs[1, :])
        z = localizer.to_spher_grid(direcs[2, :])
        #scat = ax.scatter(x, y, z, s=100)
    if EXTERNAL_PLOT:
        fig = plt.figure()
        ax = fig.add_subplot(111)
        plt.show(block=False)

    count = 0
    try:
        global done
        while in_stream.is_active() or out_stream.is_active():
            data_available = in_buf.wait_for_read(WINDOW_LENGTH, TIMEOUT)
            if data_available:
                if switch_beamforming:
                    DO_BEAMFORM = not DO_BEAMFORM
                    switch_beamforming = False
                    # Get data from the circular buffer
                data = in_buf.read_samples(WINDOW_LENGTH)
                # Perform an stft
                stft.performStft(data)
                # Process dfts from windowed segments of input
                dfts = stft.getDFTs()
                rffts = mat.to_all_real_matlab_format(dfts)
                d, energy = localizer.get_distribution_real(
                    rffts[:, :, 0], 'gcc')  # Use first hop
                post = localizer.get_distribution(rffts[:, :, 0])
                ind = np.argmax(post)
                u = 1.5 * direcs[:, ind]  # Direction of arrival
                #if energy < 500:
                #    continue

                # Do beam forming
                if DO_BEAMFORM:
                    align_mat = align_mats[:, :, ind]
                    filtered = beamformer.filter_real(rffts, align_mat)
                    mat.set_dfts_real(dfts, filtered, n_channels=2)

                # Take care of plotting
                if count % 1 == 0:
                    if PLOT_CARTES:
                        ax.cla()
                        ax.grid(False)
                        d = localizer.to_spher_grid(
                            post / (np.max(post) + consts.EPS))
                        #d = localizer.to_spher_grid(d / (np.max(d) + consts.EPS))
                        ax.scatter(x, y, z, c=d, s=40)
                        #ax.plot_surface(x, y, z, rstride=1, cstride=1, facecolor=plt.cm.gist_heat(d))
                        ax.plot([0, u[0]], [0, u[1]], [0, u[2]],
                                c='black',
                                linewidth=3)
                        if DO_BEAMFORM:
                            if np.max(np.abs(response)) > 1:
                                response /= np.max(np.abs(response))
                            X = response * x
                            Y = response * y
                            Z = response * z
                            ax.plot_surface(X,
                                            Y,
                                            Z,
                                            rstride=1,
                                            cstride=1,
                                            color='white')
                        ax.set_xlim(-1, 1)
                        ax.set_ylim(-1, 1)
                        ax.set_zlim(0, 1)
                        #ax.view_init(90, -90)
                        fig.canvas.draw()
                    if PLOT_2D:
                        # Get unconditional distribution
                        dist = localizer.to_spher_grid(d)
                        dist -= np.min(dist)
                        dist /= (np.sum(dist) + consts.EPS)
                        sample_mat[:, :-1] = sample_mat[:, 1:]
                        sample_mat[:, -1] = dist
                        # Get kalman estimate
                        maxind = np.argmax(post)
                        estimate_mat[:-1] = estimate_mat[1:]
                        estimate_mat[-1] = maxind
                        plot_2d.set_array(sample_mat)
                        state_est_plot.set_ydata(estimate_mat)
                        plt.draw()
                count += 1

                # Get the istft of the processed data
                if PLAY_AUDIO or RECORD_AUDIO:
                    new_data = stft.performIStft()
                    new_data = out_buf.reduce_channels(new_data,
                                                       NUM_CHANNELS_IN,
                                                       NUM_CHANNELS_OUT)
                    # Write out the new, altered data
                    if PLAY_AUDIO:
                        if out_buf.get_available_write() >= WINDOW_LENGTH:
                            out_buf.write_samples(new_data)
                    if RECORD_AUDIO:
                        if record_buf.get_available_write() >= WINDOW_LENGTH:
                            record_buf.write_samples(new_data)

    except KeyboardInterrupt:
        print "Program interrupted"
        done = True

    print "Cleaning up"
    in_stream.stop_stream()
    in_stream.close()
    out_stream.stop_stream()
    out_stream.close()
    pa.terminate()

    # Take care of output file
    if RECORD_AUDIO:
        print "Writing output file"
        make_wav()

    print "Done"