hydrus support

README.md

@@ -146,6 +146,25 @@ shape: (1, 512), dtype: uint8
    42 181  38 254 177 232 150  99]]
 ```
 
+# Server and hydrus compatibility
+
+This feature is only for python3.
+
+It is tested on python 3.6.5 on ubuntu  18.04.
+
+``i2v`` can run a local server by doing the following:
+
+- put `illust2vec_tag_ver200.caffemodel` and `tag_list.json` on current working command
+- to run on on `127.0.0.1` host and port `5011` run following command:
+```shell
+  $ i2v run -h 127.0.0.1 -p 5011
+```
+
+Hydrus can use that as parsing by importing following config:
+```json
+[32, "illustration2vec", 2, ["http://127.0.0.1:5011/image/new/?url=%2Fimage%2Fplausible-tag", 1, 0, [55, 1, [[], "some hash bytes"]], "path", {}, [[29, 1, ["match parser", [27, 6, [[26, 1, [[62, 2, [0, "a", {"id": "hydrus-link"}, null, null, false, [51, 1, [3, "", null, null, "example string"]]]]]], 0, "href", [51, 1, [3, "", null, null, "example string"]], [55, 1, [[], "parsed information"]]]], [[30, 3, ["tags creator", 0, [27, 6, [[26, 1, [[62, 2, [0, "td", {"class": "tag-creator"}, null, null, false, [51, 1, [3, "", null, null, "example string"]]]]]], 1, "", [51, 1, [3, "", null, null, "example string"]], [55, 1, [[], "parsed information"]]]], "creator"]], [30, 3, ["tags series", 0, [27, 6, [[26, 1, [[62, 2, [0, "td", {"class": "tag-copyright"}, null, null, false, [51, 1, [3, "", null, null, "example string"]]]]]], 1, "", [51, 1, [3, "", null, null, "example string"]], [55, 1, [[], "parsed information"]]]], "series"]], [30, 3, ["tags character", 0, [27, 6, [[26, 1, [[62, 2, [0, "td", {"class": "tag-character"}, null, null, false, [51, 1, [3, "", null, null, "example string"]]]]]], 1, "", [51, 1, [3, "", null, null, "example string"]], [55, 1, [[], "parsed information"]]]], "character"]], [30, 3, ["tags unnamespaced", 0, [27, 6, [[26, 1, [[62, 2, [0, "td", {"class": "tag-general"}, null, null, false, [51, 1, [3, "", null, null, "example string"]]]]]], 1, "", [51, 1, [3, "", null, null, "example string"]], [55, 1, [[], "parsed information"]]]], ""]], [30, 3, ["tags rating", 0, [27, 6, [[26, 1, [[62, 2, [0, "td", {"class": "tag-rating", "data-index": "1"}, null, null, false, [51, 1, [3, "", null, null, "example string"]]]]]], 1, "", [51, 1, [3, "", null, null, "example string"]], [55, 1, [[], "parsed information"]]]], "rating"]]]]]]]]
+```
+
 # License
 The pre-trained models and the other files we have provided are licensed
 under the MIT License.

i2v/__main__.py

@@ -0,0 +1,95 @@
+#!/usr/bin/env python
+from pprint import pprint
+import os
+import os.path as op
+import sys
+
+from flask import Flask, __version__ as flask_version, send_from_directory
+from flask.cli import FlaskGroup
+from flask_admin import Admin
+from PIL import Image
+import click
+
+from . import make_i2v_with_chainer, views, models
+
+
+__version__ = '0.2.1'
+
+
+class CustomFlaskGroup(FlaskGroup):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.params[0].help = 'Show the program version'
+        self.params[0].callback = get_custom_version
+
+
+def get_custom_version(ctx, param, value):
+    if not value or ctx.resilient_parsing:
+        return
+    message = 'Illustration2Vec %(app_version)s\nFlask %(version)s\nPython %(python_version)s'
+    click.echo(message % {
+        'app_version': __version__,
+        'version': flask_version,
+        'python_version': sys.version,
+    }, color=ctx.color)
+    ctx.exit()
+
+
+def create_app():
+    app = Flask(__name__)
+    app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = 'True'
+    app.config['DATABASE_FILE'] = 'i2v.sqlite'
+    app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///' + app.config['DATABASE_FILE']
+    app.config['SECRET_KEY'] = os.getenv('ILLUSTRATION2VEC_SECRET_KEY') or os.urandom(24)
+    # Create directory for file fields to use
+    try:
+        os.mkdir(models.file_path)
+        app.logger.debug('File path created')
+    except OSError:
+        pass
+    # app and db
+    models.db.init_app(app)
+    app.app_context().push()
+    models.db.create_all()
+    # other setup
+    admin = Admin(
+        app, name='Illustration2Vec', template_mode='bootstrap3',
+        index_view=views.HomeView(url='/')
+    )
+    admin.add_view(views.ImageView(models.Image, models.db.session))
+    admin.add_view(views.ChecksumView(models.Checksum, models.db.session))
+    admin.add_view(views.TagEstimationView(models.TagEstimation, models.db.session))
+    app.add_url_rule('/file/<filename>', 'file', view_func=lambda filename: send_from_directory(models.file_path, filename))
+    app.logger.debug('file path: {}'.format(models.file_path))
+    return app
+
+
+@click.group(cls=CustomFlaskGroup, create_app=create_app)
+def cli():
+    """Illustration2Vec."""
+    pass
+
+
+@cli.command()
+@click.option('--output', help='Output format;[default]/pprint', default='default')
+@click.argument('images', nargs=-1)
+def estimate_plausible_tags(images, output='default'):
+    """Estimate plausible tags."""
+    illust2vec = make_i2v_with_chainer(
+        "illust2vec_tag_ver200.caffemodel", "tag_list.json")
+    image_sets = map(lambda x: {'value': x, 'sha256': models.sha256_checksum(x)}, images)
+    for idx, img_set in enumerate(image_sets):
+        img = Image.open(img_set['value'])
+        res = illust2vec.estimate_plausible_tags([img], threshold=0.5)
+        if isinstance(res[idx]['rating'], zip):
+            res[idx]['rating'] = list(res[idx]['rating'])
+        print("image: {}\nsha256: {}".format(
+            img_set['value'], img_set['sha256']))
+        if output == 'pprint':
+            pprint(res)
+        else:
+            print(res)
+
+
+if __name__ == '__main__':
+    cli()

i2v/chainer_i2v.py

@@ -2,11 +2,16 @@
 import json
 import warnings
 import numpy as np
+from distutils.version import StrictVersion
 from scipy.ndimage import zoom
 from skimage.transform import resize
+import chainer
 from chainer import Variable
 from chainer.functions import average_pooling_2d, sigmoid
-from chainer.functions.caffe import CaffeFunction
+try:
+    from chainer.functions.caffe import CaffeFunction
+except:
+    from chainer.links.caffe import CaffeFunction
 
 
 class ChainerI2V(Illustration2VecBase):
@@ -47,7 +52,14 @@ def _forward(self, inputs, layername):
         input_ -= self.mean  # subtract mean
         input_ = input_.transpose((0, 3, 1, 2))  # (N, H, W, C) -> (N, C, H, W)
         x = Variable(input_)
-        y, = self.net(inputs={'data': x}, outputs=[layername], train=False)
+
+        # train argument is not supported from Ver2.
+        if StrictVersion(chainer.__version__) < StrictVersion('2.0.0'):
+            y, = self.net(inputs={'data': x}, outputs=[layername], train=False)
+        else:
+            chainer.using_config('train', False)
+            y, = self.net(inputs={'data': x}, outputs=[layername])
+
         return y
 
     def _extract(self, inputs, layername):

i2v/models.py

@@ -0,0 +1,168 @@
+#!/usr/bin/env python3
+"""Model module."""
+from datetime import datetime
+import hashlib
+import os
+import os.path as op
+
+from appdirs import user_data_dir
+from flask_admin import form
+from flask_sqlalchemy import SQLAlchemy
+from sqlalchemy.event import listens_for
+from sqlalchemy.types import TIMESTAMP
+from sqlalchemy_utils.types.choice import ChoiceType
+
+MODE_PLAUSIBLE_TAG = 'plausible'
+MODE_TOP_TAG = 'top'
+db = SQLAlchemy()
+file_path = op.join(user_data_dir('Illustration2Vec', 'Masaki Saito'), 'files')
+
+
+class Base(db.Model):
+    __abstract__ = True
+    id = db.Column(db.Integer, primary_key=True)
+    created_at = db.Column(TIMESTAMP, default=datetime.now, nullable=False)
+
+
+class Image(Base):
+    id = db.Column(db.Integer, primary_key=True)
+    path = db.Column(db.String)
+    checksum_id = db.Column(db.Integer, db.ForeignKey('checksum.id'))
+    checksum = db.relationship(
+        'Checksum', foreign_keys='Image.checksum_id', lazy='subquery',
+        backref=db.backref('images', lazy=True, cascade='delete'))
+
+    def __repr__(self):
+        return '<Image {0.id}>'.format(self)
+
+    @property
+    def full_path(self):
+        return os.path.join(file_path, self.path)
+
+    def update_checksum(self, session=None):
+        session = db.session if session is None else session
+        checksum_val = sha256_checksum(self.full_path)
+        self.checksum = get_or_create(session, Checksum, value=checksum_val)[0]
+
+    @property
+    def thumbgen_filename(self):
+        return form.thumbgen_filename(self.path)
+
+
+def sha256_checksum(filename, block_size=65536):
+    sha256 = hashlib.sha256()
+    with open(filename, 'rb') as f:
+        for block in iter(lambda: f.read(block_size), b''):
+            sha256.update(block)
+    return sha256.hexdigest()
+
+
+@listens_for(Image, 'after_delete')
+def del_image(mapper, connection, target):
+    if target.path:
+        # Delete image
+        try:
+            os.remove(op.join(file_path, target.path))
+        except OSError:
+            pass
+        # Delete thumbnail
+        try:
+            os.remove(op.join(file_path, form.thumbgen_filename(target.path)))
+        except OSError:
+            pass
+
+
+class Checksum(Base):
+    value = db.Column(db.String, unique=True)
+
+    def update_tag_estimation(self, tags, mode=MODE_PLAUSIBLE_TAG, session=None):
+        session = db.session if session is None else session
+        for nm, list_value in tags.items():
+            if list_value:
+                for tag_value, estimation_value in list_value:
+                    tag_model = get_or_create_tag(
+                        value=str(tag_value), namespace=nm, session=session)[0]
+                    e_item = get_or_create(
+                        session, TagEstimation,
+                        checksum=self, tag=tag_model, mode=mode)[0]
+                    e_item.value = estimation_value
+                    yield e_item
+
+    def __repr__(self):
+        return '<Checksum {0.id} {0.value}>'.format(self)
+
+    def get_estimated_tags(self, mode=MODE_PLAUSIBLE_TAG):
+        res = {'character': [], 'copyright': [], 'general': [], 'rating': []}
+        for estimation in self.tag_estimations:
+            if estimation.mode == mode:
+                res.setdefault(estimation.tag.namespace.value, []).append(
+                    (estimation.tag.value, estimation.value))
+        return res
+
+
+class TagEstimation(Base):
+    MODES = [
+        (MODE_PLAUSIBLE_TAG, MODE_PLAUSIBLE_TAG),
+        (MODE_TOP_TAG, MODE_TOP_TAG),
+    ]
+    checksum_id = db.Column(db.Integer, db.ForeignKey('checksum.id'))
+    checksum = db.relationship(
+        'Checksum', foreign_keys='TagEstimation.checksum_id', lazy='subquery',
+        backref=db.backref('tag_estimations', lazy=True))
+    tag_id = db.Column(db.Integer, db.ForeignKey('tag.id'))
+    tag = db.relationship(
+        'Tag', foreign_keys='TagEstimation.tag_id', lazy='subquery',
+        backref=db.backref('tag_estimations', lazy=True))
+    value = db.Column(db.Float)
+    mode = db.Column(ChoiceType(MODES))
+
+    def __repr__(self):
+        templ = \
+            '<TagEstimation {0.id} mode:{0.mode.value} {0.tag.fullname} confidence:{1}>'
+        return templ.format(
+            self, '{0:.2f}'.format(self.value * 100)
+        )
+
+
+def get_or_create_tag(value, namespace=None, session=None):
+    session = db.session if session is None else session
+    kwargs = dict(value=value)
+    if namespace:
+        namespace_model = get_or_create(session, Namespace, value=namespace)[0]
+        kwargs['namespace'] = namespace_model
+    model, created = get_or_create(session, Tag, **kwargs)
+    return model, created
+
+
+class Tag(Base):
+    value = db.Column(db.String)
+    namespace_id = db.Column(db.Integer, db.ForeignKey('namespace.id'))
+    namespace = db.relationship(
+        'Namespace', foreign_keys='Tag.namespace_id', lazy='subquery',
+        backref=db.backref('tags', lazy=True))
+
+    @property
+    def fullname(self):
+        res = ''
+        if self.namespace:
+            res = self.namespace.value + ':'
+        res += self.value
+        return res
+
+    def __repr__(self):
+        return '<Tag {0.id} {0.fullname}>'.format(self)
+
+
+class Namespace(Base):
+    value = db.Column(db.String, unique=True)
+
+
+def get_or_create(session, model, **kwargs):
+    """Creates an object or returns the object if exists."""
+    instance = session.query(model).filter_by(**kwargs).first()
+    created = False
+    if not instance:
+        instance = model(**kwargs)
+        session.add(instance)
+        created = True
+    return instance, created

i2v/templates/i2v/home.html

@@ -0,0 +1,7 @@
+{% extends 'admin/master.html' %}
+{% block head %}
+<link rel="shortcut icon" href="#" />
+{% endblock %}
+{% block body %}
+<p>hello world</p>
+{% endblock %}