Python3に対応させたい

gaze.py

1import os 2import argparse 3import csv 4import numpy as np 5from matplotlib import pyplot, image 6 7def draw_display(dispsize, imagefile=None): 8 """Returns a matplotlib.pyplot Figure and its axes, with a size of 9 dispsize, a black background colour, and optionally with an image drawn 10 onto it 11 12 arguments 13 14 dispsize - tuple or list indicating the size of the display, 15 e.g. (1024,768) 16 17 keyword arguments 18 19 imagefile - full path to an image file over which the heatmap 20 is to be laid, or None for no image; NOTE: the image 21 may be smaller than the display size, the function 22 assumes that the image was presented at the centre of 23 the display (default = None) 24 25 returns 26 fig, ax - matplotlib.pyplot Figure and its axes: field of zeros 27 with a size of dispsize, and an image drawn onto it 28 if an imagefile was passed 29 """ 30 31 # construct screen (black background) 32 screen = np.zeros((dispsize[1], dispsize[0], 3), dtype='float32') 33 # if an image location has been passed, draw the image 34 if imagefile != None: 35 # check if the path to the image exists 36 if not os.path.isfile(imagefile): 37 raise Exception("ERROR in draw_display: imagefile not found at '%s'" % imagefile) 38 # load image 39 img = image.imread(imagefile) 40 41 # width and height of the image 42 w, h = len(img[0]), len(img) 43 # x and y position of the image on the display 44 x = dispsize[0] // 2 - w // 2 45 y = dispsize[1] // 2 - h // 2 46 # draw the image on the screen 47 screen[y:y + h, x:x + w, :] += img 48 # dots per inch 49 dpi = 100.0 50 # determine the figure size in inches 51 figsize = (dispsize[0] // dpi, dispsize[1] // dpi) 52 # create a figure 53 fig = pyplot.figure(figsize=figsize, dpi=dpi, frameon=False) 54 ax = pyplot.Axes(fig, [0, 0, 1, 1]) 55 ax.set_axis_off() 56 fig.add_axes(ax) 57 # plot display 58 ax.axis([0, dispsize[0], 0, dispsize[1]]) 59 ax.imshow(screen) # , origin='upper') 60 61 return fig, ax 62 63def gaussian(x, sx, y=None, sy=None): 64 """Returns an array of numpy arrays (a matrix) containing values between 65 1 and 0 in a 2D Gaussian distribution 66 67 arguments 68 x -- width in pixels 69 sx -- width standard deviation 70 71 keyword argments 72 y -- height in pixels (default = x) 73 sy -- height standard deviation (default = sx) 74 """ 75 76 # square Gaussian if only x values are passed 77 if y == None: 78 y = x 79 if sy == None: 80 sy = sx 81 # centers 82 xo = x // 2 83 yo = y // 2 84 # matrix of zeros 85 M = np.zeros([y, x], dtype=float) 86 # gaussian matrix 87 for i in range(x): 88 for j in range(y): 89 M[j, i] = np.exp( 90 -1.0 * (((float(i) - xo) ** 2 // (2 * sx * sx)) + ((float(j) - yo) ** 2 // (2 * sy * sy)))) 91 92 return M 93 94def draw_heatmap(gazepoints, dispsize, imagefile=None, alpha=0.5, savefilename=None, gaussianwh=200, gaussiansd=None): 95 """Draws a heatmap of the provided fixations, optionally drawn over an 96 image, and optionally allocating more weight to fixations with a higher 97 duration. 98 99 arguments 100 101 gazepoints - a list of gazepoint tuples (x, y) 102 103 dispsize - tuple or list indicating the size of the display, 104 e.g. (1024,768) 105 106 keyword arguments 107 108 imagefile - full path to an image file over which the heatmap 109 is to be laid, or None for no image; NOTE: the image 110 may be smaller than the display size, the function 111 assumes that the image was presented at the centre of 112 the display (default = None) 113 alpha - float between 0 and 1, indicating the transparancy of 114 the heatmap, where 0 is completely transparant and 1 115 is completely untransparant (default = 0.5) 116 savefilename - full path to the file in which the heatmap should be 117 saved, or None to not save the file (default = None) 118 119 returns 120 121 fig - a matplotlib.pyplot Figure instance, containing the 122 heatmap 123 """ 124 125 # IMAGE 126 fig, ax = draw_display(dispsize, imagefile=imagefile) 127 128 # HEATMAP 129 # Gaussian 130 gwh = gaussianwh 131 gsdwh = gwh // 6 if (gaussiansd is None) else gaussiansd 132 gaus = gaussian(gwh, gsdwh) 133 # matrix of zeroes 134 strt = gwh // 2 135 heatmapsize = dispsize[1] + 2 * strt, dispsize[0] + 2 * strt 136 heatmap = np.zeros(heatmapsize, dtype=float) 137 # create heatmap 138 for i in range(0, len(gazepoints)): 139 # get x and y coordinates 140 x = strt + gazepoints[i][0] - int(gwh // 2) 141 y = strt + gazepoints[i][1] - int(gwh // 2) 142 # correct Gaussian size if either coordinate falls outside of 143 # display boundaries 144 if (not 0 < x < dispsize[0]) or (not 0 < y < dispsize[1]): 145 hadj = [0, gwh]; 146 vadj = [0, gwh] 147 if 0 > x: 148 hadj[0] = abs(x) 149 x = 0 150 elif dispsize[0] < x: 151 hadj[1] = gwh - int(x - dispsize[0]) 152 if 0 > y: 153 vadj[0] = abs(y) 154 y = 0 155 elif dispsize[1] < y: 156 vadj[1] = gwh - int(y - dispsize[1]) 157 # add adjusted Gaussian to the current heatmap 158 try: 159 heatmap[y:y + vadj[1], x:x + hadj[1]] += gaus[vadj[0]:vadj[1], hadj[0]:hadj[1]] * gazepoints[i][2] 160 except: 161 # fixation was probably outside of display 162 pass 163 else: 164 # add Gaussian to the current heatmap 165 heatmap[y:y + gwh, x:x + gwh] += gaus * gazepoints[i][2] 166 # resize heatmap 167 heatmap = heatmap[strt:dispsize[1] + strt, strt:dispsize[0] + strt] 168 # remove zeros 169 lowbound = np.mean(heatmap[heatmap > 0]) 170 heatmap[heatmap < lowbound] = np.NaN 171 # draw heatmap on top of image Greys 172 ax.imshow(heatmap, cmap='jet', alpha=alpha) 173 174 # FINISH PLOT 175 # invert the y axis, as (0,0) is top left on a display 176 ax.invert_yaxis() 177 # save the figure if a file name was provided 178 if savefilename != None: 179 fig.savefig(savefilename) 180 181 return fig 182 183 184################## 185# Parsing # 186################## 187 188parser = argparse.ArgumentParser(description='Parameters required for processing.') 189 190#required args 191parser.add_argument('input-path', type=str, help='path to the csv input') 192parser.add_argument('display-width', type=int, help='an integer representing the display width') 193parser.add_argument('display-height', type=int, help='an integer representing the display height') 194 195#optional args 196parser.add_argument('-a', '--alpha', type=float, default='0.5', required=False, help='alpha for the gaze overlay') 197parser.add_argument('-o', '--output-name', type=str, required=False, help='name for the output file') 198parser.add_argument('-b', '--background-image', type=str, default=None, required=False, help='path to the background image') 199 200#advanced optional args 201parser.add_argument('-n', '--n-gaussian-matrix', type=int, default='200', required=False, help='width and height of gaussian matrix') 202parser.add_argument('-sd', '--standard-deviation', type=float, default=None ,required=False, help='standard deviation of gaussian distribution') 203 204 205args = vars(parser.parse_args()) 206 207input_path = args['input-path'] 208display_width = args['display-width'] 209display_height = args['display-height'] 210alpha = args['alpha'] 211output_name = args['output_name'] if args['output_name'] is not None else 'output' 212background_image = args['background_image'] 213ngaussian = args['n_gaussian_matrix'] 214sd = args['standard_deviation'] 215 216with open(input_path) as f: 217 reader = csv.reader(f) 218 raw = list(reader) 219 220 gaza_data = [] 221 if len(raw[0]) == 2: 222 gaze_data = list(map(lambda q: (int(q[0]), int(q[1]), 1), raw)) 223 else: 224 gaze_data = list(map(lambda q: (int(q[0]), int(q[1]), int(q[2])), raw)) 225 226 draw_heatmap(gaze_data, (display_width, display_height), alpha=alpha, savefilename=output_name, imagefile=background_image, gaussianwh=ngaussian, gaussiansd=sd) 227

コメントを投稿

0 コメント