接受base参数，允许绘图时指定底图。

examples/test.py

@@ -9,23 +9,49 @@
 
 from pyheatmap.heatmap import HeatMap
 
-def main():
 
-    # 加载测试数据
-    sdata = open("test_data.txt").read().split("\n")
+def loadDataFromFile(fn):
+
+    lines = open(fn).read().split("\n")
     data = []
-    for ln in sdata:
+    for ln in lines:
         a = ln.split(",")
         if len(a) != 2:
             continue
         a = [int(i) for i in a]
         data.append(a)
 
+    return data
+
+
+def example2():
+
+    data_1 = loadDataFromFile("test_data.txt")
+    data_2 = loadDataFromFile("test_data2.txt")
+
+    hm = HeatMap(data_1)
+    hit_img = hm.clickmap()
+    hm2 = HeatMap(data_2)
+    hit_img2 = hm2.clickmap(base=hit_img, color=(0, 0, 255, 255))
+    hit_img2.save("hit2.png")
+
+
+def example1():
+
+    # 加载测试数据
+    data = loadDataFromFile("test_data.txt")
+
     # 开始绘制
     hm = HeatMap(data)
     hm.clickmap(save_as="hit.png")
     hm.heatmap(save_as="heat.png")
 
+
+def main():
+#    example1()
+    example2()
+
+
 if __name__ == "__main__":
     main()
 

pyheatmap/heatmap.py

@@ -44,12 +44,14 @@ def __init__(self,
             self.height = self.height or h
 
 
-    def __mkImg(self):
+    def __mkImg(self, base=None):
         u"""生成临时图片"""
 
-        if self.base:
-            self.__im = Image.open(self.base)
-            self.width, self.height = self.__im.size()
+        base = base or self.base
+
+        if base:
+            self.__im = Image.open(base) if type(base) in (str, unicode) else base
+            self.width, self.height = self.__im.size
 
         else:
             self.__im = Image.new("RGBA", (self.width, self.height), (0, 0, 0, 0))
@@ -74,10 +76,10 @@ def __paintHit(self, x, y, color):
                     im.putpixel((ix, iy), color)
 
 
-    def clickmap(self, save_as, color=(255, 0, 0, 255)):
+    def clickmap(self, save_as=None, base=None, color=(255, 0, 0, 255)):
         u"""绘制点击图片"""
 
-        self.__mkImg()
+        self.__mkImg(base)
 
         for hit in self.data:
             x, y = hit[0], hit[1]
@@ -87,8 +89,11 @@ def clickmap(self, save_as, color=(255, 0, 0, 255)):
             self.__paintHit(x, y, color)
 
 
-        self.save_as = save_as
-        self.__save()
+        if save_as:
+            self.save_as = save_as
+            self.__save()
+
+        return self.__im
 
 
     def __heat(self, heat_data, x, y, template):
@@ -137,7 +142,7 @@ def __paintHeat(self, heat_data, colors):
                     dr.point((x, y), fill=color)
 
 
-    def heatmap(self, save_as):
+    def heatmap(self, save_as=None, base=None):
         u"""绘制热图"""
 
         self.__mkImg()
@@ -154,8 +159,11 @@ def heatmap(self, save_as):
 
         self.__paintHeat(heat_data, cf.mkColors())
 
-        self.save_as = save_as
-        self.__save()
+        if save_as:
+            self.save_as = save_as
+            self.__save()
+
+        return self.__im
 
 
     def __save(self):

setup.py

@@ -4,7 +4,7 @@
 
 setup(
     name="pyheatmap",
-    version="0.1.3",
+    version="0.1.5",
     packages=["pyheatmap", "pyheatmap.inc"],
     url="https://github.com/oldj/pyheatmap",
     license="LGPL",