thesis/fourier/main.py

   1 #!/usr/bin/python
   2
   3 import sys
   4 import os
   5 from PIL import Image
   6 import Image, ImageDraw
   7 from fourier import *
   8
   9 import Gnuplot, Gnuplot.funcutils
  10
  11 gnuplot = Gnuplot.Gnuplot()
  12
  13 def levels(image, window=None):
  14         pix = image.load()
  15         data = [[255] * (image.size[1]) for _ in range(image.size[0])]
  16         for x in range(0, image.size[0]):
  17                 for y in range(0, image.size[1]):
  18
  19                         #print " Image size is " + str(image.size)
  20                         #print "x = " + str(x) + " / " + str(len(data)) + " ; y = " + str(y) + " / " + str(len(data[x]))
  21                         data[x][y] = 0.0 + pix[x,y]
  22                         if (window != None):
  23                                 data[x][y] *= window[x][y]
  24                         #print str(data[x][y])
  25         return data
  26
  27 def makePlaneWave(k, data=None, amplitude=255.0):
  28         if data == None:
  29                 data = [ [0] * 512 for _ in range(512) ]
  30         for x in range(len(data)):
  31                 for y in range(len(data[x])):
  32                         data[x][y] += 0.5 * (amplitude * cmath.exp(-1.0j * (k[0] * x + k[1] * y)) + amplitude)
  33         return data
  34
  35 def constructImages(data, image_type="fft"):
  36         images = [Image.new("L",(len(data), len(data[0])), "white") for _ in range(2)]
  37         for i in range(0,2):
  38                 pix = images[i].load()
  39                 for x in range(0, len(data)):
  40                         for y in range(0, len(data[x])):
  41                                 if i == 0:
  42                                         pix[x,y] = int(math.log(abs(data[x][y]) ** 2.0)) if image_type=="fft" else int(data[x][y].real)
  43                                 else:
  44                                         if image_type == "fft":
  45                                                 pix[x,y] = int(math.atan2(data[x][y].imag, data[x][y].real) / math.pi * 180.0)
  46                                         else:
  47                                                 pix[x,y] = int(data[x][y].imag)
  48         return images
  49
  50 def plotFile(filename, output="",title="",phase=False):
  51
  52         f = open(filename, "r")
  53         f.readline()
  54         size = f.readline().strip("\r\n# \t").split(" ")
  55         size = map(float, [size[1], size[3]])
  56         scale = f.readline().strip("\r\n# \t").split(" ")
  57         if scale[1] == "None":
  58                 scale = None
  59                 units = "pixels"
  60         else:
  61                 if (len(scale) >= 5):
  62                         units = str(scale[4])
  63                 else:
  64                         units = "pixels"
  65                 scale = map(float, [scale[1], scale[3]])
  66
  67         f.close()
  68
  69         out = open(".tmp", "w")
  70         if title == "": title = filename
  71
  72         if scale == None:
  73                 xsize = size[0]
  74                 ysize = size[1]
  75         else:
  76                 xsize = float(size[0]) * float(scale[1]) / float(scale[0])
  77                 ysize = float(size[1]) * float(scale[1]) / float(scale[0])
  78
  79
  80         if (output != ""):
  81                 out.write("set term png size 800,800\n")
  82                 out.write("set output \""+str(output)+"\"\n")
  83
  84
  85         out.write("set title \""+str(title)+"\"\n")
  86
  87         out.write("set xlabel \"Kx (1/"+str(units)+")\"\n")
  88         out.write("set ylabel \"Ky (1/"+str(units)+")\"\n")
  89         out.write("set size square\n")
  90         out.write("set key off\n")
  91
  92
  93         xtics = ""
  94         ytics = ""
  95         for tic in [-1.0, -0.75, -0.25, -0.5, 0, 0.25, 0.5, 0.75, 1.0]:
  96                 xtics = xtics + "\""+ ("%0.2f" % float(xsize * tic))+"\" "+ ("%0.2f" % float(size[0] * tic))+","
  97                 ytics = ytics + "\""+ ("%0.2f" % float(ysize * tic))+"\" "+ ("%0.2f" % float(size[1] * tic))+","
  98
  99
 100
 101         xtics = "(" + xtics[0:len(xtics)-1] + ")"
 102         ytics = "(" + ytics[0:len(ytics)-1] + ")"
 103
 104         out.write("set xtics "+str(xtics)+"\n")
 105         out.write("set ytics "+str(ytics)+"\n")
 106
 107         u = "($1 - "+str(size[0]/2)+"):($2 - "+str(size[1]/2)+")"
 108         u = u + ":(log($3))" if not phase else u + ":4"
 109
 110         out.write("set xrange [-"+str(size[0]/2)+":"+str(size[0]/2)+"]\n")
 111         out.write("set yrange [-"+str(size[1]/2)+":"+str(size[1]/2)+"]\n")
 112
 113         out.write("plot \""+str(filename)+"\" u " + str(u) + " w image")
 114
 115         out.close()
 116
 117         result = os.system("gnuplot --persist .tmp")
 118
 119         os.remove(".tmp")
 120         return result
 121
 122 def pngsave(im, file):
 123     # these can be automatically added to Image.info dict
 124     # they are not user-added metadata
 125     reserved = ('interlace', 'gamma', 'dpi', 'transparency', 'aspect')
 126
 127     # undocumented class
 128     from PIL import PngImagePlugin
 129     meta = PngImagePlugin.PngInfo()
 130
 131     # copy metadata into new object
 132     for k,v in im.info.iteritems():
 133         if k in reserved: continue
 134         meta.add_text(k, v, 0)
 135
 136     # and save
 137     im.save(file, "PNG", pnginfo=meta)
 138
 139 def writeData(data, filename, scale = None):
 140         out = open(filename, "w", 0)
 141         out.write("# FFT Datafile\n")
 142         out.write("# Size: "+str(len(data)) + " x " + str(len(data[0])) + "\n")
 143         out.write("# Scale: ")
 144         if scale != None:
 145                 out.write(str(scale[0]) + " = " + str(scale[1]) + "\n")
 146         else:
 147                 out.write("None\n")
 148
 149         for x in range(0,len(data)):
 150                 for y in range(0, len(data[x])):
 151                         out.write(str(x) + "\t" + str(y) + "\t")
 152                         out.write(str(abs(data[x][y])**2.0) + "\t")
 153                         out.write(str(math.atan2(data[x][y].imag, data[x][y].real) / math.pi * 180.0) + "\n")
 154                 out.write("\n")
 155         out.close()
 156
 157 def readData(filename):
 158         f = open(filename, "r")
 159         f.readline()
 160         size = f.readline().strip("\r\n# \t").split(" ")
 161         size = map(int, [size[1], size[3]])
 162         scale = f.readline().strip("\r\n# \t").split(" ")
 163
 164         data = [ [0] * size[1] for _ in range(size[0]) ]
 165
 166         for x in range(size[0]):
 167                 for y in range(size[1]):
 168                         s = f.readline().split("\t")
 169                         while (len(s) != 4):
 170                                 s = f.readline().strip("\r\n\t ").split("\t")
 171                         if int(s[0]) != x or int(s[1]) != y:
 172                                 sys.stderr.write("readData; bad\n")
 173                                 f.close()
 174                                 return data
 175                         mag = float(s[2])
 176                         phase = float(s[3]) * math.pi / 180.0
 177                         data[x][y] = mag * (math.cos(phase) + math.sin(phase) * 1.0j)
 178
 179         f.close()
 180         return data
 181
 182
 183 def getSlice(data, angle):
 184         x = 0
 185         y = 0
 186         result = []
 187
 188         good = True
 189         k = 0
 190         while good:
 191                 good = False
 192                 x1 = len(data)/2.0 - x
 193                 y1 = len(data)/2.0 - y
 194                 x2 = len(data)/2.0 + x
 195                 y2 = len(data)/2.0 + y
 196
 197                 if (int(x1) != int(x2) or int(y1) != int(y2)):
 198                         if (int(x1) >= 0 and int(y1) >= 0) and (int(x1) < len(data) and int(y1) < len(data)):
 199                                 result = [[-k, data[int(x1)][int(y1)]]] + result
 200                                 #print str(result) + " from " + str(x1) + ","+str(y1)
 201                                 good = True
 202
 203                 if (int(x2) >= 0 and int(y2) >= 0) and (int(x2) < len(data) and int(y2) < len(data)):
 204                         result.append([k,data[int(x2)][int(y2)]])
 205                         #print str(result) + " from " + str(x2) + ","+str(y2)
 206                         good = True
 207                 k += 1
 208                 x += math.cos(angle)
 209                 y += math.sin(angle)
 210         return result
 211
 212 def plotSlice(s):
 213         out = open(".tmp.dat", "w")
 214
 215         for line in s:
 216                 out.write(str(line[0]) + "\t" + str(abs(line[1])**2.0) + "\t" + str(math.atan2(line[1].imag, line[1].real) * math.pi / 180.0) + "\n")
 217         out.close()
 218         out = open(".tmp", "w")
 219         out.write("plot \".tmp.dat\" u 1:(log($2)) w lp\n")
 220         out.close()
 221
 222         result = os.system("gnuplot --persist .tmp")
 223
 224         #os.remove(".tmp.dat")
 225         os.remove(".tmp")
 226         return result
 227
 228 def makeWindow(w, h):
 229         result = [ [0] * h for _ in range(w)]
 230         for x in range(0, w):
 231                 for y in range(0, h):
 232                         pass
 233
 234
 235 def main(argv):
 236         s = argv[1].split(".")
 237         if s[len(s)-1] == "fft":
 238                 if (len(argv) > 2):
 239                         return plotFile(argv[1], output=argv[2])
 240                 else:
 241                         return plotFile(argv[1])
 242
 243         image = Image.open(argv[1]).convert("L")
 244         scale = map(float, image.info["scale"].split(" ")) if "scale" in image.info else None
 245
 246         data = levels(image)
 247
 248         dataf = offset2d(fft2d(data))
 249
 250
 251         if (len(argv) > 2):
 252                 writeData(dataf, argv[2], scale)
 253         else:
 254                 writeData(dataf, argv[1]+".fft", scale)
 255
 256         return 0
 257
 258
 259 if __name__ == "__main__":
 260         sys.exit(main(sys.argv))
 261
 262