TextMapsのアノテーション環境をlabelImgで構築

TextMapsは、ディープラーニングを用いた視覚ベースのウェブコンテンツ抽出ツールです。このツールのためのアノテーション構築環境をします。

labelImg

labelImgは、Pascal VOCフォーマットに対応した画像認識のためのアノテーション環境です。このアノテーション環境が使えるので、問題は以下のように定義できます。

"Wepページのスクリーンショットと、Webページ内の候補要素の矩形位置をPascal VOCで定義し、候補要素に対して要素の種類をアノテーションする。"

labelImgの使い方は以下でも見てください。 https://qiita.com/wakaba130/items/e86109b3cbd1b0dde902

TextMapsのダウンローダを使う

TextMapsではダウンローダが用意されていますが、それを用いれば、1)Webのスクリーンショット, 2)domツリーの２つを取得できます。(phantomjsが必要です。)

def download_page(url):
    print "Downloading:",url
    temp_dir = tempfile.mkdtemp()
    result = subprocess.check_output(["phantomjs", "download_page.js",url,temp_dir])
    return temp_dir


if __name__ == "__main__":
    try:
        download_dir = download_page(url)
    except subprocess.CalledProcessError:
        print "Download was not succesfull"
        sys.exit(1)

これにより、スクリーンショット(jpeg)とdomツリー(json)のファイルが生成されます。

ほしいのは、各要素のポジションなので、それを以下のような方法で取得します。

with open("tmp/dom.json") as f:
    data = eval(f.read())

def search_content(parent):
    results = []
    stack = [parent['childNodes']]
    while(True):
        children = stack.pop(0)
        for child in children:
            if "name" in child and child["name"] in ["#text", "IMG"]:
                results.append(child['position'])
            if "childNodes" in child:
                stack.append(child['childNodes'])
    return results

positions = search_content(parent)

Pascal VOCへ変換するモジュール

候補の矩形位置からPascal VOCへ変換する以下のようなモジュールを作成します。ただし、python2.7です。

import dicttoxml
import re
import lxml.etree as etree
from os.path import join


def dict_transformer(dict_data):
    regex = re.compile(r'<\?xml version="1.0" encoding="UTF-8" \?><root>(.*)</root>')
    return re.sub(regex, r'\1', dicttoxml.dicttoxml(dict_data, attr_type=False))


def dict_builder(fileid, positions, fileformat="jpeg", default_label="none"):
    objs = [{"name":default_label, "bndbox":{"xmin":x[0], "ymin":x[1], "xmax":x[2], "ymax":x[3]}} for x in positions]
    out = {"annotation":{"filename": str(fileid)+".{}".format(fileformat),"object": objs}}
    return out


def fixer(transformed_xml, remove=["object"], item_rep=("item","object")):
    c = transformed_xml[:]
    for r in remove:
        c = c.replace("<"+r+">", "").replace("</"+r+">","")
    return c.replace(item_rep[0], item_rep[1])


def prettify(xml_str):
    root = etree.fromstring(xml_str)
    return etree.tostring(root, pretty_print=True)


def annotate(fileid, positions, outpath=".", fileformat="jpeg", default_label="none"):
    result = dict_builder(fileid, positions, fileformat, default_label)
    outfile = str(fileid)+".xml"
    with open(join(outpath, outfile), "w") as f:
        funcs = [dict_transformer, fixer, prettify, f.write]
        for func in funcs:
            result = func(result)
    return result

要件は以下です。

1. id.jpegという形式でスクリーンショットを大量保存。
2. id.jsonという形式でdom treeファイルをスクリーンショットに対応させて保存。
3. annotate関数にスクリーンショットidとそのスクリーンショット内の候補ポジション一覧を渡すと、pascal voc形式のxmlファイルができる。

max_fileid = 1000
for i in range(max_fileid+1):
    with open("{}.json".format(i)) as f:
        positions = search_content(json.load(f))
    annotate(i, positions)

すると、スクリーンショットとPascal VOCファイルが1対1で対応している状態なので、これらを同じディレクトリに格納します。

mkdir example
mv *.xml example
mv *.jpeg example

あとはlabelImgでOpenDirをクリックし、exampleを開くだけです。

参考

[0] https://github.com/tzutalin/labelImg [1] https://github.com/gogartom/TextMaps