convert.py

#coding:utf-8
# Copyright (c) 2020  PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import os
import time
import tarfile
import shutil
from string import Template

from paddlehub.common import tmp_dir
from paddlehub.commands.base_command import BaseCommand, ENTRY
from paddlehub.common.hub_server import CacheUpdater

INIT_FILE = '__init__.py'
MODULE_FILE = 'module.py'
SERVING_FILE = 'serving_client_demo.py'
TMPL_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'tmpl')


class ConvertCommand(BaseCommand):
    name = "convert"

    def __init__(self, name):
        super(ConvertCommand, self).__init__(name)
        self.show_in_help = True
        self.description = "Convert model to PaddleHub-Module."
        self.parser = argparse.ArgumentParser(
            description=self.__class__.__doc__,
            prog='%s %s [COMMAND]' % (ENTRY, name),
            usage='%(prog)s',
            add_help=True)
        self.parser.add_argument('command')
        self.parser.add_argument('--module_name', '-n')
        self.parser.add_argument('--module_version',
                                 '-v',
                                 nargs='?',
                                 default='1.0.0')
        self.parser.add_argument('--model_dir', '-d')
        self.parser.add_argument('--output_dir', '-o')

    def create_module_tar(self):
        if not os.path.exists(self.dest):
            os.makedirs(self.dest)
        tar_file = os.path.join(self.dest, '{}.tar.gz'.format(self.module))
        with tarfile.open(tar_file, 'w:gz') as tfp:
            tfp.add(self.dest, recursive=False, arcname=self.module)
            for root, dir, files in os.walk(self.src):
                for file in files:
                    fullpath = os.path.join(root, file)
                    arcname = os.path.join(self.module, 'assets', file)
                    tfp.add(fullpath, arcname=arcname)

            tfp.add(self.model_file,
                    arcname=os.path.join(self.module, MODULE_FILE))
            tfp.add(self.serving_file,
                    arcname=os.path.join(self.module, SERVING_FILE))
            tfp.add(self.init_file,
                    arcname=os.path.join(self.module, INIT_FILE))

    def create_module_py(self):
        template_file = open(os.path.join(TMPL_DIR, 'x_model.tmpl'),
                             'r',
                             encoding='utf-8')
        tmpl = Template(template_file.read())
        lines = []

        lines.append(
            tmpl.substitute(NAME="'{}'".format(self.module),
                            TYPE="'CV'",
                            AUTHOR="'Baidu'",
                            SUMMARY="''",
                            VERSION="'{}'".format(self.version),
                            EMAIL="''"))
        # self.model_file = os.path.join(self.dest, MODULE_FILE)
        self.model_file = os.path.join(self._tmp_dir, MODULE_FILE)
        if os.path.exists(self.model_file):
            raise RuntimeError(
                'File `{MODULE_FILE}` is already exists in src dir.'.format(
                    MODULE_FILE))

        with open(self.model_file, 'w', encoding='utf-8') as fp:
            fp.writelines(lines)

    def create_init_py(self):
        # self.init_file = os.path.join(self.dest, INIT_FILE)
        self.init_file = os.path.join(self._tmp_dir, INIT_FILE)
        if os.path.exists(self.init_file):
            return
        shutil.copyfile(os.path.join(TMPL_DIR, 'init_py.tmpl'), self.init_file)

    def create_serving_demo_py(self):
        template_file = open(os.path.join(TMPL_DIR, 'serving_demo.tmpl'),
                             'r',
                             encoding='utf-8')
        tmpl = Template(template_file.read())
        lines = []

        lines.append(tmpl.substitute(MODULE_NAME=self.module))
        # self.serving_file = os.path.join(self.dest, SERVING_FILE)
        self.serving_file = os.path.join(self._tmp_dir, SERVING_FILE)
        if os.path.exists(self.serving_file):
            raise RuntimeError(
                'File `{}` is already exists in src dir.'.format(SERVING_FILE))

        with open(self.serving_file, 'w', encoding='utf-8') as fp:
            fp.writelines(lines)

    @staticmethod
    def show_help():
        str = "convert --module <module> [--version <version>] --dest dest_dir --src srd_dir\n"
        str += "\tConvert model to PaddleHub-Module.\n"
        str += "--model_dir\n"
        str += "\tDir of model you want to export.\n"
        str += "--module_name:\n"
        str += "\tSet name of module.\n"
        str += "--module_version\n"
        str += "\tSet version of module, default is `1.0.0`.\n"
        str += "--output_dir\n"
        str += "\tDir to save PaddleHub-Module after exporting, default is `.`.\n"
        print(str)

        return

    def execute(self, argv):
        args = self.parser.parse_args()

        if not args.module_name or not args.model_dir:
            ConvertCommand.show_help()
            return False
        self.module = args.module_name
        self.version = args.module_version if args.module_version is not None else '1.0.0'
        self.src = args.model_dir
        self.dest = args.output_dir if args.output_dir is not None else os.path.join(
            '{}_{}'.format(self.module, str(time.time())))

        CacheUpdater("hub_convert", self.module, self.version).start()
        os.makedirs(self.dest)

        with tmp_dir() as _dir:
            self._tmp_dir = _dir
            self.create_module_py()
            self.create_init_py()
            self.create_serving_demo_py()
            self.create_module_tar()

        print('The converted module is stored in `{}`.'.format(self.dest))

        return True

    def run(self, module, version, src, dest):

        self.module = module
        self.version = version
        self.src = src
        self.dest = dest

        os.makedirs(self.dest)

        with tmp_dir() as _dir:
            self._tmp_dir = _dir
            self.create_module_py()
            self.create_init_py()
            self.create_serving_demo_py()
            self.create_module_tar()

        return True


command = ConvertCommand.instance()

if __name__ == '__main__':
    command.run('test_module_name', '1.1.1', './new_model', './new_module')