+#!/usr/bin/python
+
+import sys
+import os
+from PIL import Image
+import Image, ImageDraw
+from fourier import *
+
+import Gnuplot, Gnuplot.funcutils
+
+gnuplot = Gnuplot.Gnuplot()
+
+def levels(image, window=None):
+ pix = image.load()
+ data = [[255] * (image.size[1]) for _ in range(image.size[0])]
+ for x in range(0, image.size[0]):
+ for y in range(0, image.size[1]):
+
+ #print " Image size is " + str(image.size)
+ #print "x = " + str(x) + " / " + str(len(data)) + " ; y = " + str(y) + " / " + str(len(data[x]))
+ data[x][y] = 0.0 + pix[x,y]
+ if (window != None):
+ data[x][y] *= window[x][y]
+ #print str(data[x][y])
+ return data
+
+def makePlaneWave(k, data=None, amplitude=255.0):
+ if data == None:
+ data = [ [0] * 512 for _ in range(512) ]
+ for x in range(len(data)):
+ for y in range(len(data[x])):
+ data[x][y] += 0.5 * (amplitude * cmath.exp(-1.0j * (k[0] * x + k[1] * y)) + amplitude)
+ return data
+
+def constructImages(data, image_type="fft"):
+ images = [Image.new("L",(len(data), len(data[0])), "white") for _ in range(2)]
+ for i in range(0,2):
+ pix = images[i].load()
+ for x in range(0, len(data)):
+ for y in range(0, len(data[x])):
+ if i == 0:
+ pix[x,y] = int(math.log(abs(data[x][y]) ** 2.0)) if image_type=="fft" else int(data[x][y].real)
+ else:
+ if image_type == "fft":
+ pix[x,y] = int(math.atan2(data[x][y].imag, data[x][y].real) / math.pi * 180.0)
+ else:
+ pix[x,y] = int(data[x][y].imag)
+ return images
+
+def plotFile(filename, output="",title="",phase=False):
+
+ f = open(filename, "r")
+ f.readline()
+ size = f.readline().strip("\r\n# \t").split(" ")
+ size = map(float, [size[1], size[3]])
+ scale = f.readline().strip("\r\n# \t").split(" ")
+ if scale[1] == "None":
+ scale = None
+ units = "pixels"
+ else:
+ if (len(scale) >= 5):
+ units = str(scale[4])
+ else:
+ units = "pixels"
+ scale = map(float, [scale[1], scale[3]])
+
+ f.close()
+
+ out = open(".tmp", "w")
+ if title == "": title = filename
+
+ if scale == None:
+ xsize = size[0]
+ ysize = size[1]
+ else:
+ xsize = float(size[0]) * float(scale[1]) / float(scale[0])
+ ysize = float(size[1]) * float(scale[1]) / float(scale[0])
+
+
+ if (output != ""):
+ out.write("set term png size 800,800\n")
+ out.write("set output \""+str(output)+"\"\n")
+
+
+ out.write("set title \""+str(title)+"\"\n")
+
+ out.write("set xlabel \"Kx (1/"+str(units)+")\"\n")
+ out.write("set ylabel \"Ky (1/"+str(units)+")\"\n")
+ out.write("set size square\n")
+ out.write("set key off\n")
+
+
+ xtics = ""
+ ytics = ""
+ for tic in [-1.0, -0.75, -0.25, -0.5, 0, 0.25, 0.5, 0.75, 1.0]:
+ xtics = xtics + "\""+ ("%0.2f" % float(xsize * tic))+"\" "+ ("%0.2f" % float(size[0] * tic))+","
+ ytics = ytics + "\""+ ("%0.2f" % float(ysize * tic))+"\" "+ ("%0.2f" % float(size[1] * tic))+","
+
+
+
+ xtics = "(" + xtics[0:len(xtics)-1] + ")"
+ ytics = "(" + ytics[0:len(ytics)-1] + ")"
+
+ out.write("set xtics "+str(xtics)+"\n")
+ out.write("set ytics "+str(ytics)+"\n")
+
+ u = "($1 - "+str(size[0]/2)+"):($2 - "+str(size[1]/2)+")"
+ u = u + ":(log($3))" if not phase else u + ":4"
+
+ out.write("set xrange [-"+str(size[0]/2)+":"+str(size[0]/2)+"]\n")
+ out.write("set yrange [-"+str(size[1]/2)+":"+str(size[1]/2)+"]\n")
+
+ out.write("plot \""+str(filename)+"\" u " + str(u) + " w image")
+
+ out.close()
+
+ result = os.system("gnuplot --persist .tmp")
+
+ os.remove(".tmp")
+ return result
+
+def pngsave(im, file):
+ # these can be automatically added to Image.info dict
+ # they are not user-added metadata
+ reserved = ('interlace', 'gamma', 'dpi', 'transparency', 'aspect')
+
+ # undocumented class
+ from PIL import PngImagePlugin
+ meta = PngImagePlugin.PngInfo()
+
+ # copy metadata into new object
+ for k,v in im.info.iteritems():
+ if k in reserved: continue
+ meta.add_text(k, v, 0)
+
+ # and save
+ im.save(file, "PNG", pnginfo=meta)
+
+def writeData(data, filename, scale = None):
+ out = open(filename, "w", 0)
+ out.write("# FFT Datafile\n")
+ out.write("# Size: "+str(len(data)) + " x " + str(len(data[0])) + "\n")
+ out.write("# Scale: ")
+ if scale != None:
+ out.write(str(scale[0]) + " = " + str(scale[1]) + "\n")
+ else:
+ out.write("None\n")
+
+ for x in range(0,len(data)):
+ for y in range(0, len(data[x])):
+ out.write(str(x) + "\t" + str(y) + "\t")
+ out.write(str(abs(data[x][y])**2.0) + "\t")
+ out.write(str(math.atan2(data[x][y].imag, data[x][y].real) / math.pi * 180.0) + "\n")
+ out.write("\n")
+ out.close()
+
+def readData(filename):
+ f = open(filename, "r")
+ f.readline()
+ size = f.readline().strip("\r\n# \t").split(" ")
+ size = map(int, [size[1], size[3]])
+ scale = f.readline().strip("\r\n# \t").split(" ")
+
+ data = [ [0] * size[1] for _ in range(size[0]) ]
+
+ for x in range(size[0]):
+ for y in range(size[1]):
+ s = f.readline().split("\t")
+ while (len(s) != 4):
+ s = f.readline().strip("\r\n\t ").split("\t")
+ if int(s[0]) != x or int(s[1]) != y:
+ sys.stderr.write("readData; bad\n")
+ f.close()
+ return data
+ mag = float(s[2])
+ phase = float(s[3]) * math.pi / 180.0
+ data[x][y] = mag * (math.cos(phase) + math.sin(phase) * 1.0j)
+
+ f.close()
+ return data
+
+
+def getSlice(data, angle):
+ x = 0
+ y = 0
+ result = []
+
+ good = True
+ k = 0
+ while good:
+ good = False
+ x1 = len(data)/2.0 - x
+ y1 = len(data)/2.0 - y
+ x2 = len(data)/2.0 + x
+ y2 = len(data)/2.0 + y
+
+ if (int(x1) != int(x2) or int(y1) != int(y2)):
+ if (int(x1) >= 0 and int(y1) >= 0) and (int(x1) < len(data) and int(y1) < len(data)):
+ result = [[-k, data[int(x1)][int(y1)]]] + result
+ #print str(result) + " from " + str(x1) + ","+str(y1)
+ good = True
+
+ if (int(x2) >= 0 and int(y2) >= 0) and (int(x2) < len(data) and int(y2) < len(data)):
+ result.append([k,data[int(x2)][int(y2)]])
+ #print str(result) + " from " + str(x2) + ","+str(y2)
+ good = True
+ k += 1
+ x += math.cos(angle)
+ y += math.sin(angle)
+ return result
+
+def plotSlice(s):
+ out = open(".tmp.dat", "w")
+
+ for line in s:
+ 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")
+ out.close()
+ out = open(".tmp", "w")
+ out.write("plot \".tmp.dat\" u 1:(log($2)) w lp\n")
+ out.close()
+
+ result = os.system("gnuplot --persist .tmp")
+
+ #os.remove(".tmp.dat")
+ os.remove(".tmp")
+ return result
+
+def makeWindow(w, h):
+ result = [ [0] * h for _ in range(w)]
+ for x in range(0, w):
+ for y in range(0, h):
+ pass
+
+
+def main(argv):
+ s = argv[1].split(".")
+ if s[len(s)-1] == "fft":
+ if (len(argv) > 2):
+ return plotFile(argv[1], output=argv[2])
+ else:
+ return plotFile(argv[1])
+
+ image = Image.open(argv[1]).convert("L")
+ scale = map(float, image.info["scale"].split(" ")) if "scale" in image.info else None
+
+ data = levels(image)
+
+ dataf = offset2d(fft2d(data))
+
+
+ if (len(argv) > 2):
+ writeData(dataf, argv[2], scale)
+ else:
+ writeData(dataf, argv[1]+".fft", scale)
+
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main(sys.argv))
+
+