StarRailCopilot/dev_tools/button_extract.py

import os

import numpy as np
from PIL import Image

os.chdir('../')

from module.base.button import get_color
from module.config.config import AzurLaneConfig
from module.logger import logger

MODULE_FOLDER = './module'
BUTTON_FILE = 'assets.py'
IMPORT_EXP = """
from module.base.button import Button
from module.base.template import Template

# This file is generated by module.dev_tools.asset_extract.
# Don't modified it manually.
"""
IMPORT_EXP = IMPORT_EXP.strip().split('\n') + ['']


class ImageExtractor:
    def __init__(self, module, file, config):
        """
        Args:
            module(str):
            file(str): xxx.png
            config(AzurLaneConfig):
        """
        self.module = module
        self.name = os.path.splitext(file)[0]
        self.config = config
        self.area, self.color, self.button = (), (), ()
        self.load()

    def file(self, genre=''):
        file = '%s.%s.png' % (self.name, genre) if genre else '%s.png' % self.name
        file = os.path.join(self.config.ASSETS_FOLDER, self.module, file)
        return file

    @staticmethod
    def extract(file):
        image = Image.open(file).convert('RGB')
        bbox = image.getbbox()
        mean = get_color(image=image, area=bbox)
        mean = tuple(np.rint(mean).astype(int))
        return bbox, mean

    def load(self):
        if os.path.exists(self.file()):
            self.area, self.color = self.extract(self.file())
            self.button = self.area
        if os.path.exists(self.file('AREA')):
            self.area, _ = self.extract(self.file('AREA'))
        if os.path.exists(self.file('COLOR')):
            _, self.color = self.extract(self.file('COLOR'))
        if os.path.exists(self.file('BUTTON')):
            self.button, _ = self.extract(self.file('BUTTON'))

    @property
    def expression(self):
        return '%s = Button(area=%s, color=%s, button=%s, file=\'%s\')' % (
            self.name, self.area, self.color, self.button,
            self.config.ASSETS_FOLDER + '/' + self.module + '/' + self.name + '.png')


class TemplateExtractor(ImageExtractor):
    # def __init__(self, module, file, config):
    #     """
    #     Args:
    #         module(str):
    #         file(str): xxx.png
    #         config(AzurLaneConfig):
    #     """
    #     self.module = module
    #     self.file = file
    #     self.config = config

    @property
    def expression(self):
        return '%s = Template(file=\'%s\')' % (
            self.name,
            self.config.ASSETS_FOLDER + '/' + self.module + '/' + self.name + '.png')
        # return '%s = Template(area=%s, color=%s, button=%s, file=\'%s\')' % (
        #     self.name, self.area, self.color, self.button,
        #     self.config.ASSETS_FOLDER + '/' + self.module + '/' + self.name + '.png')


# class OcrExtractor(ImageExtractor):
#     @property
#     def expression(self):
#         return '%s = OcrArea(area=%s, color=%s, button=%s, file=\'%s\')' % (
#             self.name, self.area, self.color, self.button,
#             self.config.ASSETS_FOLDER + '/' + self.module + '/' + self.name + '.png')


class ModuleExtractor:
    def __init__(self, name, config):
        self.name = name
        self.config = config
        self.folder = os.path.join(self.config.ASSETS_FOLDER, name)

    @staticmethod
    def split(file):
        name, ext = os.path.splitext(file)
        name, sub = os.path.splitext(name)
        return name, sub, ext

    def is_base_image(self, file):
        _, sub, _ = self.split(file)
        return sub == ''

    @property
    def expression(self):
        exp = []
        for file in os.listdir(self.folder):
            if file[0].isdigit():
                continue
            if file.startswith('TEMPLATE_'):
                exp.append(TemplateExtractor(module=self.name, file=file, config=self.config).expression)
                continue
            # if file.startswith('OCR_'):
            #     exp.append(OcrExtractor(module=self.name, file=file, config=self.config).expression)
            #     continue
            if self.is_base_image(file):
                exp.append(ImageExtractor(module=self.name, file=file, config=self.config).expression)
                continue

        logger.info('Module: %s(%s)' % (self.name, len(exp)))
        exp = IMPORT_EXP + exp
        return exp

    def write(self):
        folder = os.path.join(MODULE_FOLDER, self.name)
        if not os.path.exists(folder):
            os.mkdir(folder)
        with open(os.path.join(folder, BUTTON_FILE), 'w') as f:
            for text in self.expression:
                f.write(text + '\n')


class AssetExtractor:
    """
    Extract Asset to asset.py.
    All the filename of assets should be in uppercase.

    Asset name starts with digit will be ignore.
        E.g. 2020XXXX.png.
    Asset name starts with 'TEMPLATE_' will treat as template.
        E.g. TEMPLATE_AMBUSH_EVADE_SUCCESS.png
             > TEMPLATE_AMBUSH_EVADE_SUCCESS = Template(file='./assets/handler/TEMPLATE_AMBUSH_EVADE_SUCCESS.png')
    Asset name starts other will treat as button.
        E.g. GET_MISSION.png
             > Button(area=(553, 482, 727, 539), color=(93, 142, 203), button=(553, 482, 727, 539), name='GET_MISSION')
    Asset name like XXX.AREA.png, XXX.COLOR.png, XXX.BUTTON.png, will overwrite the attribute of XXX.png.
        E.g. BATTLE_STATUS_S.BUTTON.png overwrites the attribute 'button' of BATTLE_STATUS_S
    Asset name starts with 'OCR_' will be treat as button.
        E.g. OCR_EXERCISE_TIMES.png.
    """

    def __init__(self, config):
        logger.info('Assets extract')
        for module in os.listdir(config.ASSETS_FOLDER):
            if os.path.isdir(os.path.join(config.ASSETS_FOLDER, module)):
                me = ModuleExtractor(name=module, config=config)
                me.write()


ae = AssetExtractor(AzurLaneConfig())