From 5026cbde84a68e897bc0dd80f9a10424fd21c93d Mon Sep 17 00:00:00 2001 From: Yuxin Wu Date: Wed, 29 Nov 2017 03:26:40 -0800 Subject: [PATCH] fix avatar rendering --- README.md | 2 +- wechat/avatar.py | 23 +++++++++++++---------- wechat/res.py | 18 ++++++++++-------- 3 files changed, 24 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index 1e47842..595a6d1 100644 --- a/README.md +++ b/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 diff --git a/wechat/avatar.py b/wechat/avatar.py index c4a5acb..b4b33dd 100644 --- a/wechat/avatar.py +++ b/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 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 diff --git a/wechat/res.py b/wechat/res.py index 5f095da..f1c8c3a 100644 --- a/wechat/res.py +++ b/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 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 - -