""" 猫狗图片分类 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()