curve_generator/plot_functions.py at master · avardy/curve_generator · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
#!/usr/bin/env python

import numpy as np
import matplotlib.pyplot as plt
import cv2
from scipy.ndimage.morphology import distance_transform_edt
from math import sqrt
from mpl_toolkits.axes_grid1 import make_axes_locatable

# Modify font size for all matplotlib plots.
plt.rcParams.update({'font.size': 10})
plt.rcParams.update({'lines.linewidth': 0.5})

def plot_image(image, title, filename=None, contours=False, contour_levels=10, pause=False):
    X, Y = np.meshgrid(range(image.shape[1]), range(image.shape[0]))

    if pause:
        plt.clf()

    ax = plt.gca()
    im = plt.imshow(image, cmap='gray', interpolation='none')
    plt.title(title)
    if contours:
        CS = plt.contour(X, Y, image, levels=contour_levels, colors='white', alpha=1.0)
        #plt.clabel(CS, inline=1, fontsize=6, manual=True)
        plt.clabel(CS, inline=1, fontsize=10)

    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    plt.colorbar(im, cax=cax)

    if filename != None:
        plt.savefig(filename, bbox_inches='tight')

    if pause:
        plt.pause(0.01)
    else:
        plt.show()

def subplot_image(image, title, axes, contours=None):
    X, Y = np.meshgrid(range(image.shape[1]), range(image.shape[0]))

    axes.imshow(image, cmap='gray', interpolation='none')
    axes.set_title(title)
    #axes.colorbar()
    if contours != None:
        CS = axes.contour(X, Y, image, levels=contours, colors='blue', alpha=1.0)
        #print("AV: Left-click to select contour label positions.  Middle-click to end.")
        axes.clabel(CS, inline=1, colors='white', fontsize=4, manual=True)


def subplot_draw_goal(goal_x, goal_y, axes):
    ll = 20
    axes.plot((goal_x-ll, goal_x+ll), (goal_y-ll, goal_y+ll), marker='', color='red')
    axes.plot((goal_x-ll, goal_x+ll), (goal_y+ll, goal_y-ll), marker='', color='red')

def subplot_gradient(distance_image, gradient, title, axes, normalize=False):
    X, Y = np.meshgrid(range(distance_image.shape[1]), range(distance_image.shape[0]))
    U = np.copy(gradient[0])
    V = np.copy(gradient[1])
    if normalize:
        for j in range(distance_image.shape[0]):
            for i in range(distance_image.shape[1]):
                u = U[j,i]
                v = V[j,i]
                if u != 0 and v != 0:
                    U[j,i] = u / sqrt(u**2 + v**2)
                    V[j,i] = v / sqrt(u**2 + v**2)

    # For the plot below set the gradient to zero except at every Nth position.
    N = 100
    for j in range(distance_image.shape[0]):
        for i in range(distance_image.shape[1]):
            if (j - N//2) % N != 0 or (i - N//2) % N != 0:
                U[j,i] = 0
                V[j,i] = 0

    axes.set_title(title)
    axes.imshow(distance_image, interpolation='none')
    #axes.quiver(X, Y, U, V, angles='xy', minlength=0, scale=40, headwidth=7),
    axes.quiver(X, Y, U, V, angles='xy', minlength=0, scale=20, headwidth=10),
    #axes.savefig('gradient.pdf', bbox_inches='tight')
    #axes.show()

def plot_levels(distance_image):
    height = distance_image.shape[0]
    width = distance_image.shape[1]

    # Normalize
    max_dist = np.amax(distance_image)
    distance_image = distance_image / max_dist

    VISIBLE_WIDTH = 1.87
    PUCK_DIAMETER = 0.03 * sqrt(2)
    MAX_DISTANCE_FACTOR = 0.8
    pixel_size = VISIBLE_WIDTH / width
    puck_radius_norm_dist = PUCK_DIAMETER / (2*pixel_size*max_dist)

    # Number of lanes
    n = 10

    threshold_triples = []
    for l in range(n):
        lo_thresh = MAX_DISTANCE_FACTOR * l / n
        hi_thresh = MAX_DISTANCE_FACTOR * (l+1) / n

        puck_detection_thresh = lo_thresh
        if l > 0:
            puck_detection_thresh += puck_radius_norm_dist

        triple = (lo_thresh, puck_detection_thresh, hi_thresh)
        threshold_triples.append(triple)

    contour_levels = [0]
    for threshold_triple in threshold_triples:
#        contour_levels.append(threshold_triple[1])
        contour_levels.append(threshold_triple[2])

    plt.imshow(distance_image, interpolation='none')
    X, Y = np.meshgrid(range(width), range(height))
    CS = plt.contour(X, Y, distance_image, levels=contour_levels, colors='w')
    plt.clabel(CS, inline=True, fontsize=18) # in-line labels for contour lines

#    plt.savefig('levels.pdf', bbox_inches='tight')
    plt.show()