Parcourir la source

Add mask image support

Xi Ruoyao il y a 5 ans
Parent
commit
07ec9fb010
1 fichiers modifiés avec 17 ajouts et 1 suppressions
  1. 17 1
      wordfreq.py

+ 17 - 1
wordfreq.py

@@ -52,7 +52,12 @@ def main():
             help = "the path to the TTF/OTF fonts to use")
     parser.add_argument("-d", "--debug", type=bool, default=False,
             help = "if you want to dump debug information")
+    parser.add_argument("-m", "--mask",
+            help = "the path to the mask image")
+    parser.add_argument("-c", "--contour", type = int, default = 0,
+            help = "the width of contour if you want it")
     args = parser.parse_args()
+
     if not args.jobs is None:
         jieba.enable_parallel(args.jobs)
     if args.input is None:
@@ -63,9 +68,20 @@ def main():
     c = cnt(words, args.parts)
     if not args.top is None:
         c = dict(c.most_common(args.top))
+
     if args.debug:
         print(c)
-    wc = WordCloud(font_path = args.fontpath).generate_from_frequencies(c)
+
+    mask = None
+    if not args.mask is None:
+        from PIL import Image
+        import numpy as np
+        mask = np.array(Image.open(args.mask))
+
+    wc = WordCloud(font_path = args.fontpath,
+            mask = mask,
+            contour_width = args.contour,
+            contour_color = "steelblue").generate_from_frequencies(c)
     wc.to_image().save(args.output)
 
 if __name__ == "__main__":