修复 base 图片的一些问题。

examples/test.py

@@ -7,6 +7,7 @@
 
 import os
 import sys
+
 module_path = os.sep.join(os.path.abspath(__file__).split(os.sep)[:-2])
 sys.path.append(module_path)
 sys.path.append(os.path.join(module_path, "pyheatmap"))
@@ -15,7 +16,6 @@
 
 
 def load_data_from_file(fn):
-
     lines = open(fn).read().split("\n")
     data = []
     for ln in lines:
@@ -29,7 +29,6 @@ def load_data_from_file(fn):
 
 
 def example2():
-
     data_1 = load_data_from_file("test_data.txt")
     data_2 = load_data_from_file("test_data2.txt")
 
@@ -41,14 +40,14 @@ def example2():
 
 
 def example1():
-
     # 加载测试数据
     data = load_data_from_file("test_data.txt")
 
     # 开始绘制
     hm = HeatMap(data)
     hm.clickmap(save_as="hit.png")
     # hm.heatmap(save_as="heat.png", r=20) # 这儿可以传入 r 参数，指定热图半径，默认 r=10
+    # hm.heatmap(save_as="heat.png", base="/var/tmp/test_base.png")
     hm.heatmap(save_as="heat.png")
 
 

pyheatmap/heatmap.py

@@ -18,6 +18,7 @@
 from .inc import cf
 
 import sys
+
 if sys.version > '3':
     PY3 = True
 else:
@@ -67,14 +68,14 @@ def __mk_img(self, base=None):
         u"""生成临时图片"""
 
         base = base or self.base
+        self.__im0 = None
 
         if base:
             str_type = (str,) if PY3 else (str, unicode)
-            self.__im = Image.open(base) if type(base) in str_type else base
-            self.width, self.height = self.__im.size
+            self.__im0 = Image.open(base) if type(base) in str_type else base
+            self.width, self.height = self.__im0.size
 
-        else:
-            self.__im = Image.new("RGBA", (self.width, self.height), (0, 0, 0, 0))
+        self.__im = Image.new("RGBA", (self.width, self.height), (0, 0, 0, 0))
 
     def __paint_hit(self, x, y, color):
         u"""绘制点击小叉图片"""
@@ -107,6 +108,7 @@ def clickmap(self, save_as=None, base=None, color=(255, 0, 0, 255), data=None):
 
             self.__paint_hit(x, y, color)
 
+        self.__add_base()
         if save_as:
             self.save_as = save_as
             self.__save()
@@ -158,6 +160,13 @@ def __paint_heat(self, heat_data, colors):
                 else:
                     dr.point((x, y), fill=color)
 
+    def __add_base(self):
+        if not self.__im0:
+            return
+
+        self.__im0.paste(self.__im, mask=self.__im)
+        self.__im = self.__im0
+
     def sample(self, max_count=None, rate=None):
 
         count = self.count
@@ -194,7 +203,7 @@ def sample(self, max_count=None, rate=None):
     def heatmap(self, save_as=None, base=None, data=None, r=10):
         u"""绘制热图"""
 
-        self.__mk_img()
+        self.__mk_img(base)
 
         circle = cf.mk_circle(r, self.width)
         heat_data = [0] * self.width * self.height
@@ -209,6 +218,7 @@ def heatmap(self, save_as=None, base=None, data=None, r=10):
             self.__heat(heat_data, x, y, n, circle)
 
         self.__paint_heat(heat_data, cf.mk_colors())
+        self.__add_base()
 
         if save_as:
             self.save_as = save_as