6 import Image, ImageDraw
9 import Gnuplot, Gnuplot.funcutils
11 gnuplot = Gnuplot.Gnuplot()
13 def levels(image, window=None):
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]):
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]
23 data[x][y] *= window[x][y]
24 #print str(data[x][y])
27 def makePlaneWave(k, data=None, amplitude=255.0):
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)
35 def constructImages(data, image_type="fft"):
36 images = [Image.new("L",(len(data), len(data[0])), "white") for _ in range(2)]
38 pix = images[i].load()
39 for x in range(0, len(data)):
40 for y in range(0, len(data[x])):
42 pix[x,y] = int(math.log(abs(data[x][y]) ** 2.0)) if image_type=="fft" else int(data[x][y].real)
44 if image_type == "fft":
45 pix[x,y] = int(math.atan2(data[x][y].imag, data[x][y].real) / math.pi * 180.0)
47 pix[x,y] = int(data[x][y].imag)
50 def plotFile(filename, output="",title="",phase=False):
52 f = open(filename, "r")
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":
65 scale = map(float, [scale[1], scale[3]])
69 out = open(".tmp", "w")
70 if title == "": title = filename
76 xsize = float(size[0]) * float(scale[1]) / float(scale[0])
77 ysize = float(size[1]) * float(scale[1]) / float(scale[0])
81 out.write("set term png size 800,800\n")
82 out.write("set output \""+str(output)+"\"\n")
85 out.write("set title \""+str(title)+"\"\n")
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")
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))+","
101 xtics = "(" + xtics[0:len(xtics)-1] + ")"
102 ytics = "(" + ytics[0:len(ytics)-1] + ")"
104 out.write("set xtics "+str(xtics)+"\n")
105 out.write("set ytics "+str(ytics)+"\n")
107 u = "($1 - "+str(size[0]/2)+"):($2 - "+str(size[1]/2)+")"
108 u = u + ":(log($3))" if not phase else u + ":4"
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")
113 out.write("plot \""+str(filename)+"\" u " + str(u) + " w image")
117 result = os.system("gnuplot --persist .tmp")
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')
128 from PIL import PngImagePlugin
129 meta = PngImagePlugin.PngInfo()
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)
137 im.save(file, "PNG", pnginfo=meta)
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: ")
145 out.write(str(scale[0]) + " = " + str(scale[1]) + "\n")
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")
157 def readData(filename):
158 f = open(filename, "r")
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(" ")
164 data = [ [0] * size[1] for _ in range(size[0]) ]
166 for x in range(size[0]):
167 for y in range(size[1]):
168 s = f.readline().split("\t")
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")
176 phase = float(s[3]) * math.pi / 180.0
177 data[x][y] = mag * (math.cos(phase) + math.sin(phase) * 1.0j)
183 def getSlice(data, angle):
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
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)
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)
213 out = open(".tmp.dat", "w")
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")
218 out = open(".tmp", "w")
219 out.write("plot \".tmp.dat\" u 1:(log($2)) w lp\n")
222 result = os.system("gnuplot --persist .tmp")
224 #os.remove(".tmp.dat")
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):
236 s = argv[1].split(".")
237 if s[len(s)-1] == "fft":
239 return plotFile(argv[1], output=argv[2])
241 return plotFile(argv[1])
243 image = Image.open(argv[1]).convert("L")
244 scale = map(float, image.info["scale"].split(" ")) if "scale" in image.info else None
248 dataf = offset2d(fft2d(data))
252 writeData(dataf, argv[2], scale)
254 writeData(dataf, argv[1]+".fft", scale)
259 if __name__ == "__main__":
260 sys.exit(main(sys.argv))