tasks/image_classification/gradio.py

"""
猫狗图片分类 Gradio 交互界面。

这个文件把图片分类 Pipeline 包装成可体验的页面，方便上传图片或选择内置示例后查看模型预测的
最大概率类别。
"""

from typing import cast

import gradio as gr
import numpy as np
from PIL import Image

from deep_learning.env.resolve import display_path, resolve_path, resolve_saved
from deep_learning.pipeline.services.checkpoint import describe_checkpoint_lookup, resolve_checkpoint
from deep_learning.pipeline.services.model_loader import load_inference_artifact
from tasks.image_classification.runner import pipeline


IMAGE_SIZE = (180, 180)
CLASS_NAMES = ["Cat", "Dog"]


class ImageClassificationTool:
    def __init__(
        self,
        pipeline,
        inference_artifact,
        resource,
        image_size: tuple[int, int],
        class_names: list[str]
    ):
        self.pipeline = pipeline
        self.inference_artifact = inference_artifact
        self.resource = resource
        self.image_size = image_size
        self.class_names = class_names

    def classify(self, image) -> str:
        input_image = Image.fromarray(np.asarray(image, dtype="uint8")).convert("RGB")
        resized_image = input_image.resize(
            (self.image_size[1], self.image_size[0])
        )
        image_array = np.asarray(resized_image, dtype="float32")
        model_input = np.expand_dims(image_array, axis=0)
        prediction = self.inference_artifact.model(model_input, training=False)
        positive_probability = float(np.asarray(prediction).reshape(-1)[0])
        class_index = int(positive_probability >= 0.5)
        probability = positive_probability
        if class_index == 0:
            probability = 1.0 - positive_probability
        return f"预测类别：{self.class_names[class_index]}，概率：{probability:.2%}"


class ImageClassificationPageAdapter:
    def __init__(self, pipeline):
        self.pipeline = pipeline
        self.sample_dir = resolve_path("data/dev/cats_vs_dogs/images")

    def list_sample_images(self) -> list[str]:
        return sorted(
            path.name
            for path in self.sample_dir.iterdir()
            if path.is_file() and path.suffix.lower() in [".jpg", ".jpeg", ".png"]
        )

    def load_sample_image(self, sample_name: str):
        image_path = self.sample_dir / sample_name
        return np.asarray(Image.open(image_path).convert("RGB"))

    def extra_model_info(self) -> str:
        return f"**示例图片目录**: {display_path(self.sample_dir)}"

    def create_tool(self, inference_artifact, resource) -> ImageClassificationTool:
        return ImageClassificationTool(
            self.pipeline,
            inference_artifact,
            resource,
            IMAGE_SIZE,
            CLASS_NAMES
        )


class ImageClassificationAppBuilder:
    def __init__(
        self,
        pipeline,
        tool_factory,
        sample_images: list[str],
        load_sample_image,
        title: str = "猫狗图片分类",
        input_label: str = "输入图片",
        sample_label: str = "示例图片",
        output_label: str = "预测结果",
        extra_model_info=None
    ):
        self.pipeline = pipeline
        self.tool_factory = tool_factory
        self.sample_images = sample_images
        self.load_sample_image = load_sample_image
        self.title = title
        self.input_label = input_label
        self.sample_label = sample_label
        self.output_label = output_label
        self.extra_model_info = extra_model_info
        self._tool = None

    def _resolve_checkpoint_rule(self) -> dict:
        checkpoint_rule = self.pipeline.checkpoint_load_rules.resolve_test_rule(
            default_dirs=[resolve_saved(f"models/{self.pipeline.name}")]
        )
        if checkpoint_rule["suffix"] is None:
            checkpoint_rule["suffix"] = ".keras"
        return checkpoint_rule

    def _load_inference_artifact(self) -> tuple:
        checkpoint_rule = self._resolve_checkpoint_rule()
        return load_inference_artifact(
            checkpoint_rule,
            self.pipeline.model_builder.load_inference_artifact
        )

    def get_model_info(self) -> str:
        parts = []
        checkpoint_rule = self._resolve_checkpoint_rule()
        checkpoint_path, _ = resolve_checkpoint(**checkpoint_rule)

        if checkpoint_path is None:
            lookup_info = describe_checkpoint_lookup(
                dirs=checkpoint_rule.get("dirs"),
                path=checkpoint_rule.get("path"),
                suffix=checkpoint_rule.get("suffix")
            )
            raise FileNotFoundError(f"未找到模型检查点文件。查找信息: {lookup_info}")

        file_name = checkpoint_path.name
        file_size = checkpoint_path.stat().st_size
        parts.append(f"**模型文件**: {file_name}（{file_size / (1024 * 1024):.2f} MB）")

        if self.extra_model_info is not None:
            extra_info = self.extra_model_info()
            if extra_info:
                parts.append(extra_info)

        return "，".join(parts)

    def _init_tool(self):
        print("正在加载图片分类模型...")
        inference_artifact, resource = self._load_inference_artifact()
        print("模型加载完成！")
        return self.tool_factory(inference_artifact, resource)

    def _ensure_tool_initialized(self) -> None:
        if self._tool is None:
            self._tool = self._init_tool()

    def select_sample_image(self, sample_name: str):
        return self.load_sample_image(sample_name)

    def classify_image(self, image):
        self._ensure_tool_initialized()
        return self._tool.classify(image)

    def create_ui(self):
        with gr.Blocks(title=self.title) as demo:
            gr.Markdown(f"# {self.title}")
            gr.Markdown("请选择示例图片或直接上传图片，分类时只会读取当前输入框中的图片。")
            try:
                gr.Markdown(self.get_model_info())
            except Exception as e:
                gr.Markdown(f"**模型信息加载失败**: {str(e)}")

            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(
                        label=self.input_label,
                        type="numpy"
                    )
                    sample_name = gr.Dropdown(
                        choices=self.sample_images,
                        value=cast(str | None, None),
                        label=self.sample_label
                    )
                    classify_btn = gr.Button("开始分类", variant="primary")

                with gr.Column():
                    output_text = gr.Textbox(
                        label=self.output_label,
                        interactive=False
                    )

            sample_name.change(
                fn=self.select_sample_image,
                inputs=[sample_name],
                outputs=[input_image]
            )

            classify_btn.click(
                fn=self.classify_image,
                inputs=[input_image],
                outputs=[output_text]
            )

        original_get_config_file = demo.get_config_file

        def get_config_file():
            config = original_get_config_file()
            for component in config["components"]:
                if component["type"] == "dropdown" and component["props"].get("label") == self.sample_label:
                    component["props"]["value"] = None
            return config

        demo.get_config_file = get_config_file
        return demo


adapter = ImageClassificationPageAdapter(pipeline)

app = ImageClassificationAppBuilder(
    pipeline=pipeline,
    tool_factory=adapter.create_tool,
    sample_images=adapter.list_sample_images(),
    load_sample_image=adapter.load_sample_image,
    extra_model_info=adapter.extra_model_info
)

demo = app.create_ui()

if __name__ == "__main__":
    from deep_learning.env.keras import enable_mixed_precision

    enable_mixed_precision()
    demo.launch()