app.py

import random
from copy import deepcopy
from pathlib import Path

import gradio as gr

from utils import KanaData, Recognizer


class App:
    def __init__(
        self,
        kana_data_path="data/kana-data.json",
        kana_char_dir="data/images/kana-chars",
        bg_image_path="data/images/bg.png",
        model_xml_path="model/model.xml",
        model_char_path="model/chars.txt",
        default_kana="あ",
        font_name="Kiwi Maru",
        css_path="style.css",
        favicon_path="favicon.png",
        brush_color="#111",
        brush_size=15,
        storage_key="kana-write-storage-key",
        storage_secret="kana-write-secret",
    ):
        self.brush_color = brush_color
        self.brush_size = brush_size
        self.bg_image_path = bg_image_path
        self.recognizer = Recognizer(model_xml_path, model_char_path)

        self.kana_data = KanaData.load(kana_data_path)
        self.kana_set = {kana for kana in self.kana_data.spell if len(kana) == 1}
        self.hira_set = {kana for v in self.kana_data.hiragana.values() for kana in v}
        self.kata_set = {kana for v in self.kana_data.katakana.values() for kana in v}

        self.kana_char_dir = Path(kana_char_dir)
        self.kana_images = [str(p) for p in self.kana_char_dir.rglob("*")]

        self.font_name = font_name
        self.css_path = css_path
        self.favicon_path = favicon_path

        self.default_kana = default_kana
        self.default_kana_image = str(self.kana_char_dir / f"{self.default_kana}.png")
        roma = self.conv_kana_to_roma(self.default_kana)
        self.default_roma = f"平假名 {self.default_kana} ({roma})"

        self.storage_key = storage_key
        self.storage_secret = storage_secret

        self.init_app()

    def is_hiragana(self, kana):
        return kana in self.hira_set

    def is_katakana(self, kana):
        return kana in self.kata_set

    def init_app(self):
        with self.init_blocks() as self.app:
            self.init_states()
            self.init_layout()
            self.init_events()
            self.init_storage()

    def init_blocks(self):
        return gr.Blocks(title="假名手寫練習")

    def init_layout(self):
        gr.HTML("<h1>✍️ 假名手寫練習</h1>")
        self.init_practice_tab()

        with gr.Sidebar(open=False):
            self.init_setting_tab()

    def init_states(self):
        self.curr_kana = gr.State(self.default_kana)
        self.curr_kana_list = gr.State(list())
        self.curr_kana_image = gr.State(self.default_kana_image)

    def init_practice_tab(self):
        with gr.Group():
            self.sketchpad = gr.Sketchpad(
                self.default_kana_image,
                type="numpy",
                image_mode="RGB",
                brush=gr.Brush(self.brush_size, self.brush_color),
                eraser=False,
                layers=False,
                label="🖊️ 寫字板",
            )

            self.target_txt = gr.Textbox(self.default_roma, lines=2, label="🎯 練習目標")
            self.result_txt = gr.TextArea(lines=2, label="💯 辨識結果")

        with gr.Row():
            self.next_btn = gr.Button("👉 下一個字")
            self.recog_btn = gr.Button("🔎 手寫辨識")

    def init_setting_tab(self):
        with gr.Accordion("⚙️ 練習設定"):
            self.use_hiragana = gr.Checkbox(True, label="練習平假名")
            self.use_katakana = gr.Checkbox(True, label="練習片假名")
            self.use_assist_chk = gr.Checkbox(True, label="顯示輔助字")
            self.use_kana_hint_chk = gr.Checkbox(True, label="練習目標提示假名")

    def init_events(self):
        recog_inputs = [self.sketchpad, self.curr_kana]
        recog_kwargs = gr_kwargs(self.recognize, recog_inputs, self.result_txt, "minimal")
        clear_kwargs = gr_kwargs(self.clear, self.curr_kana_image, self.sketchpad)

        next_inputs = [self.use_hiragana, self.use_katakana, self.use_assist_chk]
        next_inputs += [self.use_kana_hint_chk, self.curr_kana_list]
        next_outputs = [self.curr_kana, self.sketchpad, self.curr_kana_image]
        next_outputs += [self.target_txt, self.result_txt, self.curr_kana_list]
        next_kwargs = gr_kwargs(self.get_rand_kana, next_inputs, next_outputs)

        update_inputs = [self.curr_kana, self.use_hiragana, self.use_katakana]
        update_inputs += [self.use_assist_chk, self.use_kana_hint_chk]
        update_inputs += [self.curr_kana_list]
        update_outputs = [self.curr_kana, self.sketchpad, self.curr_kana_image]
        update_outputs += [self.target_txt, self.curr_kana_list]
        update_kwargs = gr_kwargs(self.update, update_inputs, update_outputs)

        self.recog_btn.click(**recog_kwargs)
        self.sketchpad.clear(**clear_kwargs)
        self.next_btn.click(**next_kwargs)
        self.use_hiragana.change(**update_kwargs)
        self.use_katakana.change(**update_kwargs)
        self.use_assist_chk.change(**update_kwargs)
        self.use_kana_hint_chk.change(**update_kwargs)

    def init_storage(self):
        components = [self.use_assist_chk, self.use_kana_hint_chk]
        triggers = [component.change for component in components]

        browser_state = gr.BrowserState(
            [component.value for component in components],
            storage_key=self.storage_key,
            secret=self.storage_secret,
        )

        self.app.load(inputs=browser_state, outputs=components)(lambda data: data)
        gr.on(triggers, inputs=components, outputs=browser_state)(lambda *data: data)

    def launch_app(self):
        font = gr.themes.GoogleFont(self.font_name)
        text_size = gr.themes.sizes.text_lg
        theme = gr.themes.Ocean(font=font, text_size=text_size)
        self.app.launch(
            theme=theme,
            css_paths=self.css_path,
            footer_links=[None],
            favicon_path=self.favicon_path,
        )

    def conv_kana_to_roma(self, kana):
        return self.kana_data.spell[kana][0]

    def init_kana_list(self):
        curr_kana_list = deepcopy(self.kana_images)
        random.shuffle(curr_kana_list)
        return curr_kana_list

    def get_kana(
        self,
        kana: str,
        use_hira: bool,
        use_kata: bool,
        use_assist: bool,
        use_kana_hint: bool,
        kana_list: list,
    ):
        kana_list = kana_list if kana_list else self.init_kana_list()
        kana_image = self.kana_char_dir / f"{kana}.png" if kana else kana_list.pop()
        kana = Path(kana_image).stem

        if use_hira ^ use_kata:
            while self.is_hiragana(kana) and not use_hira:
                kana_list = kana_list if kana_list else self.init_kana_list()
                kana_image = kana_list.pop()
                kana = Path(kana_image).stem
            while self.is_katakana(kana) and not use_kata:
                kana_list = kana_list if kana_list else self.init_kana_list()
                kana_image = kana_list.pop()
                kana = Path(kana_image).stem

        kana_image = kana_image if use_assist else self.bg_image_path

        kana_type = "平假名" if self.is_hiragana(kana) else "片假名"
        roma = self.conv_kana_to_roma(kana)
        roma = f"{kana_type} {kana} ({roma})" if use_kana_hint else f"{kana_type} {roma}"

        return kana, str(kana_image), roma, kana_list

    def get_rand_kana(self, use_hira, use_kata, use_assist, use_hint, kana_list):
        args = (None, use_hira, use_kata, use_assist, use_hint, kana_list)
        kana, image, roma, kana_list = self.get_kana(*args)
        return kana, image, image, roma, None, kana_list

    def update(self, kana, use_hira, use_kata, use_assist, use_hint, kana_list):
        args = (kana, use_hira, use_kata, use_assist, use_hint, kana_list)
        kana, image, roma, kana_list = self.get_kana(*args)
        return kana, image, image, roma, kana_list

    def recognize(self, image, curr_kana):
        image = image["layers"][0]
        image[image == 0] = 255
        image[image != 255] = 0
        _, results = self.recognizer.recognize(image)

        return f"正解：{curr_kana}\n辨識：" + ", ".join(
            f"{result.char} ({self.conv_kana_to_roma(result.char)}) {result.prob:.2%}"
            for items in results
            for result in items
            if result.prob > 1e-2 and result.char in self.kana_set
        )

    def clear(self, curr_kana_image):
        return curr_kana_image


def gr_kwargs(fn, inputs=None, outputs=None, show_progress="hidden", **kwargs):
    return dict(
        fn=fn,
        inputs=inputs,
        outputs=outputs,
        show_progress=show_progress,
        **kwargs,
    )


if __name__ == "__main__":
    App().launch_app()