StarRailCopilot/dev_tools/button_extract.py

import os

import numpy as np
from PIL import Image

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') + ['']
VALID_SERVER = ['cn', 'en']


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.file = {}, {}, {}, {}
        for server in VALID_SERVER:
            self.load(server)

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

    @staticmethod
    def extract(file):
        image = Image.open(file).convert('RGB')
        if image.size != (1280, 720):
            logger.warning(f'Incorrect asset {image.size} {file}')
        bbox = image.getbbox()
        mean = get_color(image=image, area=bbox)
        mean = tuple(np.rint(mean).astype(int))
        return bbox, mean

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

            self.area[server] = area
            self.color[server] = color
            self.button[server] = button
            self.file[server] = f"{self.config.ASSETS_FOLDER}/{server}/{self.module}/{self.name}.png"
        else:
            logger.attr(server, f'{self.name} not found, use cn server assets')
            self.area[server] = self.area['cn']
            self.color[server] = self.color['cn']
            self.button[server] = self.button['cn']
            self.file[server] = self.file['cn']

    @property
    def expression(self):
        return '%s = Button(area=%s, color=%s, button=%s, file=%s)' % (
            self.name, self.area, self.color, self.button, self.file)


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
    @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

    @property
    def expression(self):
        return '%s = Template(file=%s)' % (
            self.name, self.file)
        # 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, 'cn', 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 + '/cn'):
            if os.path.isdir(os.path.join(config.ASSETS_FOLDER + '/cn', module)):
                me = ModuleExtractor(name=module, config=config)
                me.write()

if __name__ == '__main__':
    ae = AssetExtractor(AzurLaneConfig('template'))