AutoAndroidController/py/venv/Lib/site-packages/modelscope/trainers/cv/cartoon_translation_trainer.py

# Copyright (c) Alibaba, Inc. and its affiliates.

import os
import os.path as osp
from typing import Dict, Optional

import numpy as np
import tensorflow as tf
from packaging import version
from tqdm import tqdm

from modelscope.models.cv.cartoon import (CartoonModel, all_file,
                                          simple_superpixel, tf_data_loader,
                                          write_batch_image)
from modelscope.trainers.base import BaseTrainer
from modelscope.trainers.builder import TRAINERS
from modelscope.utils.constant import ModelFile
from modelscope.utils.logger import get_logger

logger = get_logger()

if version.parse(tf.__version__) < version.parse('2'):
    pass
else:
    logger.info(
        f'TensorFlow version {_tf_version} found, TF2.x is not supported by CartoonTranslationTrainer.'
    )


@TRAINERS.register_module(module_name=r'cartoon-translation')
class CartoonTranslationTrainer(BaseTrainer):

    def __init__(self,
                 model: str,
                 cfg_file: str = None,
                 work_dir=None,
                 photo=None,
                 cartoon=None,
                 max_steps=None,
                 *args,
                 **kwargs):
        """
                Args:
                    model: the model_id of trained model
                    cfg_file: the path of configuration file
                    work_dir: the path to save training results
                    photo: the path of photo images for training
                    cartoon: the path of cartoon images for training
                    max_steps: the number of total iteration for training
                Returns:
                    initialized trainer: object of CartoonTranslationTrainer
        """
        model = self.get_or_download_model_dir(model)
        tf.reset_default_graph()

        self.model_dir = model
        self.model_path = osp.join(model, ModelFile.TF_CHECKPOINT_FOLDER)
        if cfg_file is None:
            cfg_file = osp.join(model, ModelFile.CONFIGURATION)

        super().__init__(cfg_file)

        self.params = {}
        self._override_params_from_file()
        if work_dir is not None:
            self.params['work_dir'] = work_dir
        if photo is not None:
            self.params['photo'] = photo
        if cartoon is not None:
            self.params['cartoon'] = cartoon
        if max_steps is not None:
            self.params['max_steps'] = max_steps

        if not os.path.exists(self.params['work_dir']):
            os.makedirs(self.params['work_dir'])

        self.face_photo_list = all_file(self.params['photo'])
        self.face_cartoon_list = all_file(self.params['cartoon'])

        tf_config = tf.ConfigProto(allow_soft_placement=True)
        tf_config.gpu_options.allow_growth = True
        self._session = tf.Session(config=tf_config)

        self.input_photo = tf.placeholder(tf.float32, [
            self.params['batch_size'], self.params['patch_size'],
            self.params['patch_size'], 3
        ])
        self.input_superpixel = tf.placeholder(tf.float32, [
            self.params['batch_size'], self.params['patch_size'],
            self.params['patch_size'], 3
        ])
        self.input_cartoon = tf.placeholder(tf.float32, [
            self.params['batch_size'], self.params['patch_size'],
            self.params['patch_size'], 3
        ])

        self.model = CartoonModel(self.model_dir)
        output = self.model(self.input_photo, self.input_cartoon,
                            self.input_superpixel)
        self.output_cartoon = output['output_cartoon']
        self.g_loss = output['g_loss']
        self.d_loss = output['d_loss']

        tf.summary.scalar('g_loss', self.g_loss)
        tf.summary.scalar('d_loss', self.d_loss)

        self.train_writer = tf.summary.FileWriter(self.params['work_dir']
                                                  + '/train_log')
        self.summary_op = tf.summary.merge_all()

        all_vars = tf.trainable_variables()
        gene_vars = [var for var in all_vars if 'gene' in var.name]
        disc_vars = [var for var in all_vars if 'disc' in var.name]
        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        with tf.control_dependencies(update_ops):
            self.g_optim = tf.train.AdamOptimizer(self.params['adv_train_lr'], beta1=0.5, beta2=0.99) \
                .minimize(self.g_loss, var_list=gene_vars)
            self.d_optim = tf.train.AdamOptimizer(self.params['adv_train_lr'], beta1=0.5, beta2=0.99) \
                .minimize(self.d_loss, var_list=disc_vars)

        self.saver = tf.train.Saver(max_to_keep=1000)
        with self._session.as_default() as sess:
            sess.run(tf.global_variables_initializer())
            if self.params['resume_epoch'] != 0:
                logger.info(f'loading model from {self.model_path}')
                self.saver.restore(
                    sess,
                    osp.join(self.model_path,
                             'model-' + str(self.params['resume_epoch'])))

    def _override_params_from_file(self):

        self.params['photo'] = self.cfg['train']['photo']
        self.params['cartoon'] = self.cfg['train']['cartoon']
        self.params['patch_size'] = self.cfg['train']['patch_size']
        self.params['work_dir'] = self.cfg['train']['work_dir']
        self.params['batch_size'] = self.cfg['train']['batch_size']
        self.params['adv_train_lr'] = self.cfg['train']['adv_train_lr']
        self.params['max_steps'] = self.cfg['train']['max_steps']
        self.params['logging_interval'] = self.cfg['train']['logging_interval']
        self.params['ckpt_period_interval'] = self.cfg['train'][
            'ckpt_period_interval']
        self.params['resume_epoch'] = self.cfg['train']['resume_epoch']
        self.params['num_gpus'] = self.cfg['train']['num_gpus']

    def train(self, *args, **kwargs):
        logger.info('Begin local cartoon translator training')

        photo_ds = tf_data_loader(self.face_photo_list,
                                  self.params['batch_size'])
        cartoon_ds = tf_data_loader(self.face_cartoon_list,
                                    self.params['batch_size'])
        photo_iterator = photo_ds.make_initializable_iterator()
        cartoon_iterator = cartoon_ds.make_initializable_iterator()
        photo_next = photo_iterator.get_next()
        cartoon_next = cartoon_iterator.get_next()

        device = 'gpu:0' if tf.test.is_gpu_available else 'cpu:0'
        with tf.device(device):

            for max_steps in tqdm(range(self.params['max_steps'])):

                self._session.run(photo_iterator.initializer)
                self._session.run(cartoon_iterator.initializer)

                photo_batch, cartoon_batch = self._session.run(
                    [photo_next, cartoon_next])

                transfer_res = self._session.run(
                    self.output_cartoon,
                    feed_dict={self.input_photo: photo_batch})

                input_superpixel = simple_superpixel(transfer_res, seg_num=200)
                g_loss, _ = self._session.run(
                    [self.g_loss, self.g_optim],
                    feed_dict={
                        self.input_photo: photo_batch,
                        self.input_superpixel: input_superpixel,
                        self.input_cartoon: cartoon_batch
                    })

                d_loss, _, train_info = self._session.run(
                    [self.d_loss, self.d_optim, self.summary_op],
                    feed_dict={
                        self.input_photo: photo_batch,
                        self.input_superpixel: input_superpixel,
                        self.input_cartoon: cartoon_batch
                    })

                self.train_writer.add_summary(train_info, max_steps)

                if np.mod(max_steps + 1, self.params['logging_interval']
                          ) == 0 or max_steps == 0:

                    logger.info(
                        f'Iter: {max_steps}, d_loss: {d_loss}, g_loss: {g_loss}'
                    )

                    if np.mod(max_steps + 1,
                              self.params['ckpt_period_interval']
                              ) == 0 or max_steps == 0:
                        self.saver.save(
                            self._session,
                            self.params['work_dir'] + '/saved_models/model',
                            write_meta_graph=False,
                            global_step=max_steps)

                        result_face = self._session.run(
                            self.output_cartoon,
                            feed_dict={
                                self.input_photo: photo_batch,
                                self.input_superpixel: photo_batch,
                                self.input_cartoon: cartoon_batch
                            })

                        write_batch_image(
                            result_face, self.params['work_dir'] + '/images',
                            str('%8d' % max_steps) + '_face_result.jpg', 4)
                        write_batch_image(
                            photo_batch, self.params['work_dir'] + '/images',
                            str('%8d' % max_steps) + '_face_photo.jpg', 4)

    def evaluate(self,
                 checkpoint_path: Optional[str] = None,
                 *args,
                 **kwargs) -> Dict[str, float]:
        """evaluate a dataset

        evaluate a dataset via a specific model from the `checkpoint_path` path, if the `checkpoint_path`
        does not exist, read from the config file.

        Args:
            checkpoint_path (Optional[str], optional): the model path. Defaults to None.

        Returns:
            Dict[str, float]: the results about the evaluation
            Example:
            {"accuracy": 0.5091743119266054, "f1": 0.673780487804878}
        """
        raise NotImplementedError(
            'evaluate is not supported by CartoonTranslationTrainer')