fix avatar rendering

README.md

@@ -84,7 +84,7 @@ pip install --user --pre pysox
 	+ What you'll need in the end is a `resource` directory with the following subdir: `avatar,emoji,image2,sfs,video,voice2`.
 
 + (Optional) Download the emoji cache from [here](https://github.com/ppwwyyxx/wechat-dump/releases/download/0.1/emoji.cache.tar.bz2)
-	and put it under `resource/emoji`. This will avoid downloading too many emojis during rendering.
+	and decompress it under `wechat-dump`. This will avoid downloading too many emojis during rendering.
 
         wget -c https://github.com/ppwwyyxx/wechat-dump/releases/download/0.1/emoji.cache.tar.bz2
         tar xf emoji.cache.tar.bz2

wechat/avatar.py

@@ -1,11 +1,12 @@
 #!/usr/bin/env python2
 # -*- coding: UTF-8 -*-
 # File: avatar.py
-# Date: Wed Nov 29 01:33:56 2017 -0800
+# Date: Wed Nov 29 03:26:10 2017 -0800
 # Author: Yuxin Wu <[email protected]>
 
 from PIL import Image
 import cStringIO
+import glob
 import os
 import numpy as np
 import logging
@@ -16,17 +17,19 @@ from common.textutil import ensure_bin_str, md5
 
 
 class AvatarReader(object):
-    def __init__(self, avt_dir, avt_db="avatar.index"):
-        self.avt_dir = avt_dir
+    def __init__(self, res_dir, avt_db="avatar.index"):
+        self.sfs_dir = os.path.join(res_dir, 'sfs')
 
         # new location of avatar, see #50
-        self.avt_dir_new = avt_dir[:avt_dir.find('sfs')] + 'avatar'
+        self.avt_dir = os.path.join(res_dir, 'avatar')
         self.avt_db = avt_db
         self._use_avt = True
-        if self.avt_db is not None and os.path.isfile(self.avt_db):
-            self.ava_use_db = True
-        elif os.path.isdir(self.avt_dir_new):
-            self.ava_use_db = False
+        if os.path.isdir(self.avt_dir) and len(os.listdir(self.avt_dir)):
+            self.avt_use_db = False
+        elif self.avt_db is not None \
+                and os.path.isfile(self.avt_db) \
+                and glob.glob(os.path.join(self.sfs_dir, 'avatar*')):
+            self.avt_use_db = True
         else:
             logger.warn(
                     "Avatar database {} not found. Will not use avatar!".format(avt_db))
@@ -49,7 +52,7 @@ class AvatarReader(object):
                     pos, size = self.query_index(filename)
                     return self.read_img(pos, size)
                 else:
-                    img_file = os.path.join(self.avt_dir_new, filename)
+                    img_file = os.path.join(self.avt_dir, filename)
                     if os.path.exists(img_file):
                         return Image.open(img_file)
                     else:
@@ -66,7 +69,7 @@ class AvatarReader(object):
 
     def read_img(self, pos, size):
         file_idx = pos >> 32
-        fname = os.path.join(self.avt_dir,
+        fname = os.path.join(self.sfs_dir,
                 'avatar.block.' + '{:05d}'.format(file_idx))
         # a 64-byte offset of each block file
         start_pos = pos - file_idx * (2**32) + 64

wechat/res.py

@@ -1,13 +1,12 @@
 #!/usr/bin/env python2
 # -*- coding: UTF-8 -*-
 # File: res.py
-# Date: Thu Jun 18 00:02:21 2015 +0800
+# Date: Wed Nov 29 03:24:17 2017 -0800
 # Author: Yuxin Wu <[email protected]>
 
 import glob
 import os
 import re
-# TODO: perhaps we don't need to introduce PIL and numpy. libjpeg might be enough
 from PIL import Image
 import cStringIO
 import base64
@@ -30,7 +29,6 @@ INTERNAL_EMOJI_DIR = os.path.join(LIB_PATH, 'static', 'internal_emoji')
 VOICE_DIRNAME = 'voice2'
 IMG_DIRNAME = 'image2'
 EMOJI_DIRNAME = 'emoji'
-AVATAR_DIRNAME = 'sfs'
 
 JPEG_QUALITY = 50
 
@@ -42,6 +40,8 @@ class EmojiCache(object):
         else:
             self.dic = {}
 
+        self._curr_size = len(self.dic)
+
     def query(self, md5):
         return self.dic.get(md5, (None, None))
 
@@ -53,7 +53,10 @@ class EmojiCache(object):
             format = im.format.lower()
             ret = (base64.b64encode(r), format)
             self.dic[md5] = ret
-            self.flush()
+
+            if len(self.dic) == self._curr_size + 10:
+                self._curr_size = len(self.dic)
+                self.flush()
             return ret
         except Exception as e:
             logger.exception("Error processing emoji from {}".format(url))
@@ -69,7 +72,7 @@ class Resource(object):
         def check(subdir):
             assert os.path.isdir(os.path.join(res_dir, subdir)), \
                     "No such directory: {}".format(subdir)
-        [check(k) for k in ['', AVATAR_DIRNAME, IMG_DIRNAME, EMOJI_DIRNAME, VOICE_DIRNAME]]
+        [check(k) for k in ['', IMG_DIRNAME, EMOJI_DIRNAME, VOICE_DIRNAME]]
 
         self.emoji_cache = EmojiCache(
                 os.path.join(os.path.dirname(os.path.abspath(__file__)),
@@ -80,7 +83,7 @@ class Resource(object):
         self.img_dir = os.path.join(res_dir, IMG_DIRNAME)
         self.voice_dir = os.path.join(res_dir, VOICE_DIRNAME)
         self.emoji_dir = os.path.join(res_dir, EMOJI_DIRNAME)
-        self.avt_reader = AvatarReader(os.path.join(res_dir, AVATAR_DIRNAME), avt_db)
+        self.avt_reader = AvatarReader(res_dir, avt_db)
 
     def get_voice_filename(self, imgpath):
         fname = md5(imgpath)
@@ -224,6 +227,7 @@ class Resource(object):
     def get_emoji_by_md5(self, md5):
         """ :returns: (b64 img, format)"""
         if md5 in self.parser.internal_emojis:
+            # TODO this seems broken
             emoji_img, format = self._get_internal_emoji(self.parser.internal_emojis[md5])
             logger.warn("Cannot get emoji {}".format(md5))
             return None, None
@@ -243,5 +247,3 @@ class Resource(object):
 
             logger.warn("Cannot get emoji {} in {}".format(md5, group))
             return None, None
-
-