Update app.py

app.py

@@ -1,15 +1,127 @@
 import os
 import json
+import warnings
+
 import gradio as gr
 import huggingface_hub
 import numpy as np
 import onnxruntime as rt
 import pandas as pd
-from PIL import Image
+from PIL import Image, ImageFile
 from huggingface_hub import login
 
 from translator import translate_texts
 
+# ------------------------------------------------------------------
+# 安全配置
+# ------------------------------------------------------------------
+# 1) 限制上传文件原始体积，拦截伪装图片/图片中塞入额外数据/高熵噪声导致的超大文件
+MAX_UPLOAD_BYTES = 8 * 1024 * 1024  # 8 MB
+# 2) 限制单边尺寸，避免异常超大分辨率
+MAX_IMAGE_SIDE = 4096
+# 3) 限制总像素数，防止“像素炸弹”或解码后内存占用过高
+MAX_IMAGE_PIXELS = 20_000_000  # 2000 万像素
+# 4) 限制解码后的估算内存占用
+MAX_DECOMPRESSED_BYTES = 160 * 1024 * 1024  # 160 MB
+# 5) 仅允许常见安全图片格式
+ALLOWED_IMAGE_FORMATS = {"PNG", "JPEG", "WEBP", "BMP", "GIF"}
+
+# Pillow 安全设置
+Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
+ImageFile.LOAD_TRUNCATED_IMAGES = False
+warnings.simplefilter("error", Image.DecompressionBombWarning)
+
+
+class ImageValidationError(ValueError):
+    """上传图片校验失败。"""
+
+
+def _format_size(num_bytes: int) -> str:
+    if num_bytes < 1024:
+        return f"{num_bytes} B"
+    if num_bytes < 1024 * 1024:
+        return f"{num_bytes / 1024:.2f} KB"
+    return f"{num_bytes / (1024 * 1024):.2f} MB"
+
+
+def validate_and_open_image(image_path: str) -> Image.Image:
+    """
+    安全打开用户上传图片：
+    - 校验原始文件体积
+    - 校验图片格式
+    - 校验宽高/总像素
+    - 校验解码后预估内存占用
+    - 拦截 Pillow 解压炸弹警告
+    """
+    if not image_path:
+        raise ImageValidationError("未检测到上传文件。")
+
+    if not os.path.isfile(image_path):
+        raise ImageValidationError("上传文件不存在或无法访问。")
+
+    file_size = os.path.getsize(image_path)
+    if file_size <= 0:
+        raise ImageValidationError("上传文件为空。")
+
+    if file_size > MAX_UPLOAD_BYTES:
+        raise ImageValidationError(
+            f"图片文件过大：{_format_size(file_size)}，超过限制 {_format_size(MAX_UPLOAD_BYTES)}。"
+        )
+
+    try:
+        with Image.open(image_path) as probe:
+            img_format = (probe.format or "").upper()
+            width, height = probe.size
+            probe.verify()
+    except Image.DecompressionBombWarning:
+        raise ImageValidationError("图片疑似像素炸弹，已被拒绝处理。")
+    except Exception as e:
+        raise ImageValidationError(f"无法解析为有效图片文件：{e}")
+
+    if img_format not in ALLOWED_IMAGE_FORMATS:
+        raise ImageValidationError(
+            f"不支持的图片格式：{img_format or '未知'}。仅允许：{', '.join(sorted(ALLOWED_IMAGE_FORMATS))}。"
+        )
+
+    if width <= 0 or height <= 0:
+        raise ImageValidationError("图片尺寸非法。")
+
+    if width > MAX_IMAGE_SIDE or height > MAX_IMAGE_SIDE:
+        raise ImageValidationError(
+            f"图片尺寸过大：{width}×{height}，单边不得超过 {MAX_IMAGE_SIDE} 像素。"
+        )
+
+    total_pixels = width * height
+    if total_pixels > MAX_IMAGE_PIXELS:
+        raise ImageValidationError(
+            f"图片总像素过大：{total_pixels:,}，超过限制 {MAX_IMAGE_PIXELS:,}。"
+        )
+
+    # 估算解码为 RGB 后的内存占用
+    estimated_decompressed_bytes = total_pixels * 3
+    if estimated_decompressed_bytes > MAX_DECOMPRESSED_BYTES:
+        raise ImageValidationError(
+            "图片解码后的内存占用过高，已拒绝处理。"
+            f" 预计占用约 {_format_size(estimated_decompressed_bytes)}，"
+            f"超过限制 {_format_size(MAX_DECOMPRESSED_BYTES)}。"
+        )
+
+    # 第二次打开，真正加载像素数据
+    try:
+        with Image.open(image_path) as img:
+            img.load()
+            if img.mode != "RGB":
+                img = img.convert("RGB")
+            else:
+                img = img.copy()
+    except Image.DecompressionBombWarning:
+        raise ImageValidationError("图片在解码阶段触发像素炸弹保护，已拒绝处理。")
+    except Exception as e:
+        raise ImageValidationError(f"图片加载失败：{e}")
+
+    return img
+
+
 # ------------------------------------------------------------------
 # 模型配置
 # ------------------------------------------------------------------
@@ -96,24 +208,23 @@ class Tagger:
                 res["characters"][tag_name] = float(outputs[idx])
                 tag_categories_for_translation["characters"].append(tag_name)
 
-
         res["general"] = dict(sorted(res["general"].items(), key=lambda kv: kv[1], reverse=True))
         res["characters"] = dict(sorted(res["characters"].items(), key=lambda kv: kv[1], reverse=True))
         res["ratings"] = dict(sorted(res["ratings"].items(), key=lambda kv: kv[1], reverse=True))
 
-
         tag_categories_for_translation["general"] = list(res["general"].keys())
         tag_categories_for_translation["characters"] = list(res["characters"].keys())
         tag_categories_for_translation["ratings"] = list(res["ratings"].keys())
 
         return res, tag_categories_for_translation
 
+
 # 全局 Tagger 实例
 try:
     tagger_instance = Tagger()
 except RuntimeError as e:
     print(f"应用启动时Tagger初始化失败: {e}")
-    tagger_instance = None # 允许应用启动，但在处理时会失败
+    tagger_instance = None  # 允许应用启动，但在处理时会失败
 
 # ------------------------------------------------------------------
 # Gradio UI
@@ -143,7 +254,7 @@ custom_css = """
 .tag-en {
     font-weight: bold;
     color: #333;
-    cursor: pointer; /* Indicates clickable */
+    cursor: pointer;
 }
 .tag-zh {
     color: #666;
@@ -153,7 +264,7 @@ custom_css = """
     color: #999;
     font-size: 0.9em;
 }
-.btn-analyze-container { /* Custom class for analyze button container */
+.btn-analyze-container {
     margin-top: 15px;
     margin-bottom: 15px;
 }
@@ -164,20 +275,17 @@ function copyToClipboard(text) {
     console.log('copyToClipboard function was called.');
     console.log('Received text:', text);
 
-    // 如果 text 未定义或为 null
     if (typeof text === 'undefined' || text === null) {
         console.warn('copyToClipboard was called with undefined or null text. Aborting this specific copy operation.');
         return;
     }
 
     navigator.clipboard.writeText(text).then(() => {
-        // console.log('Tag copied to clipboard: ' + text);
         const feedback = document.createElement('div');
-        
-        // 确保 text 是字符串类型，再进行 substring 操作
+
         let displayText = String(text);
         displayText = displayText.substring(0, 30) + (displayText.length > 30 ? '...' : '');
-        
+
         feedback.textContent = '已复制: ' + displayText;
         feedback.style.position = 'fixed';
         feedback.style.bottom = '20px';
@@ -193,7 +301,7 @@ function copyToClipboard(text) {
         setTimeout(() => {
             feedback.style.opacity = '0';
             setTimeout(() => {
-                if (document.body.contains(feedback)) { // 确保元素还在DOM中
+                if (document.body.contains(feedback)) {
                     document.body.removeChild(feedback);
                 }
             }, 500);
@@ -227,24 +335,27 @@ function copyToClipboard(text) {
 
 with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=custom_css, js=_js_functions) as demo:
     gr.Markdown("# 🖼️ AI 图像标签分析器")
-    gr.Markdown("上传图片自动识别标签，支持中英文显示和一键复制。[NovelAI在线绘画](https://nai.idlecloud.cc/)")
+    gr.Markdown(
+        "上传图片自动识别标签，支持中英文显示和一键复制。"
+        "[NovelAI在线绘画](https://nai.idlecloud.cc/)\n\n"
+    )
 
     state_res = gr.State({})
     state_translations_dict = gr.State({})
     state_tag_categories_for_translation = gr.State({})
 
-
     with gr.Row():
         with gr.Column(scale=1):
-            img_in = gr.Image(type="pil", label="上传图片", height=300)
-            
+            # 改为 filepath，确保可以拿到原始文件路径与体积进行校验
+            img_in = gr.Image(type="filepath", label="上传图片", height=300)
+
             btn = gr.Button("🚀 开始分析", variant="primary", elem_classes=["btn-analyze-container"])
-            
+
             with gr.Accordion("⚙️ 高级设置", open=False):
                 gen_slider = gr.Slider(0, 1, value=0.35, step=0.01, label="通用标签阈值", info="越高 → 标签更少更准")
                 char_slider = gr.Slider(0, 1, value=0.85, step=0.01, label="角色标签阈值", info="推荐保持较高阈值")
                 show_tag_scores = gr.Checkbox(True, label="在列表中显示标签置信度")
-            
+
             with gr.Accordion("📊 标签汇总设置", open=True):
                 gr.Markdown("选择要包含在下方汇总文本框中的标签类别：")
                 with gr.Row():
@@ -265,7 +376,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
                     out_char = gr.HTML(label="Character Tags")
                 with gr.TabItem("⭐ 评分标签"):
                     out_rating = gr.HTML(label="Rating Tags")
-            
+
             gr.Markdown("### 标签汇总结果")
             out_summary = gr.Textbox(
                 label="标签汇总",
@@ -278,9 +389,9 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
     def format_tags_html(tags_dict, translations_list, category_name, show_scores=True, show_translation_in_list=True):
         if not tags_dict:
             return "<p>暂无标签</p>"
-        
+
         html = '<div class="label-container">'
-        
+
         if not isinstance(translations_list, list):
             translations_list = []
 
@@ -288,14 +399,14 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
 
         for i, tag in enumerate(tag_keys):
             score = tags_dict[tag]
-            escaped_tag = tag.replace("'", "\\'") # Escape for JS
-            
+            escaped_tag = tag.replace("'", "\\'")
+
             html += '<div class="tag-item">'
             tag_display_html = f'<span class="tag-en" onclick="copyToClipboard(\'{escaped_tag}\')">{tag}</span>'
-            
+
             if show_translation_in_list and i < len(translations_list) and translations_list[i]:
                 tag_display_html += f'<span class="tag-zh">({translations_list[i]})</span>'
-            
+
             html += f'<div>{tag_display_html}</div>'
             if show_scores:
                 html += f'<span class="tag-score">{score:.3f}</span>'
@@ -315,9 +426,12 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
         separator = separators.get(s_sep_type, ", ")
 
         categories_to_summarize = []
-        if s_gen: categories_to_summarize.append("general")
-        if s_char: categories_to_summarize.append("characters")
-        if s_rat: categories_to_summarize.append("ratings")
+        if s_gen:
+            categories_to_summarize.append("general")
+        if s_char:
+            categories_to_summarize.append("characters")
+        if s_rat:
+            categories_to_summarize.append("ratings")
 
         if not categories_to_summarize:
             return "请至少选择一个标签类别进行汇总。"
@@ -333,83 +447,99 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
                         tags_to_join.append(f"{en_tag}/*{cat_translations[i]}*/")
                     else:
                         tags_to_join.append(en_tag)
-                if tags_to_join: # only add if there are tags for this category
-                     summary_parts.append(separator.join(tags_to_join))
-        
+                if tags_to_join:
+                    summary_parts.append(separator.join(tags_to_join))
+
         joiner = "\n\n" if separator != "\n" and len(summary_parts) > 1 else separator if separator == "\n" else " "
-        
+
         final_summary = joiner.join(summary_parts)
         return final_summary if final_summary else "选定的类别中没有找到标签。"
 
-
     def process_image_and_generate_outputs(
-        img, g_th, c_th, s_scores, # Main inputs
+        image_path, g_th, c_th, s_scores,
         s_gen, s_char, s_rat, s_sep, s_zh_in_sum
-        ):
-        if img is None:
+    ):
+        if image_path is None:
             yield (
                 gr.update(interactive=True, value="🚀 开始分析"),
                 gr.update(visible=True, value="❌ 请先上传图片。"),
                 "", "", "", "",
-                gr.update(placeholder="请先上传图片并开始分析..."),
-                {}, {}, {} 
+                {}, {}, {}
             )
             return
-        
+
         if tagger_instance is None:
             yield (
                 gr.update(interactive=True, value="🚀 开始分析"),
                 gr.update(visible=True, value="❌ 分析器未成功初始化，请检查控制台错误。"),
                 "", "", "", "",
-                gr.update(placeholder="分析器初始化失败..."),
                 {}, {}, {}
             )
             return
 
         yield (
             gr.update(interactive=False, value="🔄 处理中..."),
-            gr.update(visible=True, value="🔄 正在分析图像，请稍候..."),
-            gr.HTML(value="<p>分析中...</p>"), # General
-            gr.HTML(value="<p>分析中...</p>"), # Character
-            gr.HTML(value="<p>分析中...</p>"), # Rating
-            gr.update(value="分析中，请稍候..."), # Summary
-            {}, {}, {} # Clear states initially
+            gr.update(visible=True, value="🔄 正在校验并分析图像，请稍候..."),
+            gr.HTML(value="<p>分析中...</p>"),
+            gr.HTML(value="<p>分析中...</p>"),
+            gr.HTML(value="<p>分析中...</p>"),
+            gr.update(value="分析中，请稍候..."),
+            {}, {}, {}
         )
-        
+
         try:
+            img = validate_and_open_image(image_path)
             res, tag_categories_original_order = tagger_instance.predict(img, g_th, c_th)
 
             all_tags_to_translate = []
             for cat_key in ["general", "characters", "ratings"]:
                 all_tags_to_translate.extend(tag_categories_original_order.get(cat_key, []))
-            
+
             all_translations_flat = []
             if all_tags_to_translate:
                 all_translations_flat = translate_texts(all_tags_to_translate, src_lang="auto", tgt_lang="zh")
-            
+
             current_translations_dict = {}
             offset = 0
             for cat_key in ["general", "characters", "ratings"]:
                 cat_original_tags = tag_categories_original_order.get(cat_key, [])
                 num_tags_in_cat = len(cat_original_tags)
                 if num_tags_in_cat > 0:
-                    current_translations_dict[cat_key] = all_translations_flat[offset : offset + num_tags_in_cat]
+                    current_translations_dict[cat_key] = all_translations_flat[offset: offset + num_tags_in_cat]
                     offset += num_tags_in_cat
                 else:
                     current_translations_dict[cat_key] = []
-            
-            general_html = format_tags_html(res.get("general", {}), current_translations_dict.get("general", []), "general", s_scores, True)
-            char_html = format_tags_html(res.get("characters", {}), current_translations_dict.get("characters", []), "characters", s_scores, True)
-            rating_html = format_tags_html(res.get("ratings", {}), current_translations_dict.get("ratings", []), "ratings", s_scores, True)
+
+            general_html = format_tags_html(
+                res.get("general", {}),
+                current_translations_dict.get("general", []),
+                "general",
+                s_scores,
+                True,
+            )
+            char_html = format_tags_html(
+                res.get("characters", {}),
+                current_translations_dict.get("characters", []),
+                "characters",
+                s_scores,
+                True,
+            )
+            rating_html = format_tags_html(
+                res.get("ratings", {}),
+                current_translations_dict.get("ratings", []),
+                "ratings",
+                s_scores,
+                True,
+            )
 
             summary_text = generate_summary_text_content(
                 res, current_translations_dict,
-                s_gen, s_char, s_rat, s_sep, s_zh_in_sum 
+                s_gen, s_char, s_rat, s_sep, s_zh_in_sum
             )
 
             yield (
                 gr.update(interactive=True, value="🚀 开始分析"),
-                gr.update(visible=True, value="✅ 分析完成!"), 
+                gr.update(visible=True, value="✅ 分析完成!"),
                 general_html,
                 char_html,
                 rating_html,
@@ -418,7 +548,17 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
                 current_translations_dict,
                 tag_categories_original_order
             )
-            
+
+        except ImageValidationError as e:
+            yield (
+                gr.update(interactive=True, value="🚀 开始分析"),
+                gr.update(visible=True, value=f"❌ 上传图片未通过安全校验：{str(e)}"),
+                "<p>图片已被安全策略拒绝</p>",
+                "<p>图片已被安全策略拒绝</p>",
+                "<p>图片已被安全策略拒绝</p>",
+                gr.update(value=f"错误: {str(e)}", placeholder="上传图片未通过安全校验..."),
+                {}, {}, {}
+            )
         except Exception as e:
             import traceback
             tb_str = traceback.format_exc()
@@ -469,4 +609,4 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
 if __name__ == "__main__":
     if tagger_instance is None:
         print("CRITICAL: Tagger 未能初始化，应用功能将受限。请检查之前的错误信息。")
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)