diff --git a/.gitignore b/.gitignore index 7e7bca2..7f320c4 100644 --- a/.gitignore +++ b/.gitignore @@ -72,4 +72,5 @@ htmlcov/ officefile data/1法律 +data/data_backup diff --git a/backend/scripts/test_pdf_image_extract.py b/backend/scripts/test_pdf_image_extract.py new file mode 100644 index 0000000..3fb28a4 --- /dev/null +++ b/backend/scripts/test_pdf_image_extract.py @@ -0,0 +1,140 @@ +"""测试 Heron 版面检测 + 裁剪 + VLM 图片描述流程""" +import argparse +import asyncio +import logging +import sys +import time +from pathlib import Path + +sys.path.append(str(Path(__file__).parent.parent)) + +from PIL import Image +from src.rag.document_loaders import PDFImageExtractor + +logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s") +logger = logging.getLogger(__name__) + +PROJECT_ROOT = Path(__file__).parent.parent.parent # course-agent-od/ +TEST_PDF = ( + PROJECT_ROOT + / "data" / "uploads" / "testuser" / "knowledge_bases" / "植物知识图谱" + / "ab810d4a-b030-45c0-94d2-bd6e2a23053b.pdf" +) +OUTPUT_DIR = PROJECT_ROOT / "data" / "images" / "_test_pdf_extract" + + +def test_extract(max_pages: int | None = None): + """阶段1：Heron 版面检测 + 裁剪""" + if not TEST_PDF.exists(): + logger.error(f"测试PDF不存在: {TEST_PDF}") + sys.exit(1) + + # 清理旧输出 + if OUTPUT_DIR.exists(): + for f in OUTPUT_DIR.iterdir(): + f.unlink() + + print(f"\n{'='*60}") + print(f"[阶段1] Heron 版面检测 + 裁剪测试") + print(f" PDF: {TEST_PDF.name}") + print(f" 输出: {OUTPUT_DIR}") + print(f"{'='*60}\n") + + t0 = time.time() + images = PDFImageExtractor.extract_images( + str(TEST_PDF), str(OUTPUT_DIR), + ) + elapsed = time.time() - t0 + + if not images: + print(" ❌ 未提取到任何图片") + return [] + + if max_pages: + images = [img for img in images if img["page"] <= max_pages] + + # 按页分组统计 + from collections import defaultdict + by_page = defaultdict(list) + for img in images: + by_page[img["page"]].append(img) + + total_size = 0 + fig_count = 0 + full_count = 0 + for page_num in sorted(by_page.keys()): + page_imgs = by_page[page_num] + for img in page_imgs: + path = Path(img["path"]) + with Image.open(path) as pil: + w, h = pil.size + size_kb = img["size"] / 1024 + total_size += img["size"] + is_full = "_full." in img["filename"] + tag = "FULL" if is_full else "FIG" + if is_full: + full_count += 1 + else: + fig_count += 1 + print( + f" page {img['page']:>3d} [{tag}] {img['filename']:<25s} " + f"{w}x{h} {size_kb:>7.1f} KB" + ) + + print( + f"\n 汇总: {len(images)} 张 ({fig_count} 裁剪 + {full_count} 整页兜底), " + f"{total_size/1024/1024:.1f} MB, 耗时 {elapsed:.1f}s\n" + ) + return images + + +async def test_vlm(images: list[dict], max_images: int = 3): + """阶段2：VLM图片描述""" + from src.llm.siliconflow import get_llm_client + + # 优先选裁剪图 + fig_images = [img for img in images if "_full." not in img["filename"]] + targets = (fig_images or images)[:max_images] + + print(f"\n{'='*60}") + print(f"[阶段2] VLM图片描述测试 ({len(targets)} 张)") + print(f"{'='*60}\n") + + client = get_llm_client() + sem = asyncio.Semaphore(3) + + async def describe_one(idx: int, img: dict): + async with sem: + t0 = time.time() + desc = await client.describe_image(img["path"], img.get("context_text", "")) + elapsed = time.time() - t0 + return idx, desc, elapsed + + tasks = [describe_one(i, img) for i, img in enumerate(targets)] + results = await asyncio.gather(*tasks, return_exceptions=True) + + for r in results: + if isinstance(r, Exception): + print(f" ❌ 失败: {r}\n") + continue + idx, desc, elapsed = r + img = targets[idx] + print(f" --- page {img['page']} ({img['filename']}) {elapsed:.1f}s ---") + print(f" {desc}\n") + + +def main(): + parser = argparse.ArgumentParser(description="测试PDF图片提取") + parser.add_argument("--vlm", action="store_true", help="启用VLM图片描述") + parser.add_argument("--max-pages", type=int, default=None, help="限制提取页数") + parser.add_argument("--max-vlm", type=int, default=3, help="VLM描述最大图片数") + args = parser.parse_args() + + images = test_extract(max_pages=args.max_pages) + + if args.vlm and images: + asyncio.run(test_vlm(images, max_images=args.max_vlm)) + + +if __name__ == "__main__": + main() diff --git a/backend/src/api/knowledge_base.py b/backend/src/api/knowledge_base.py index 378f3d0..ef66098 100644 --- a/backend/src/api/knowledge_base.py +++ b/backend/src/api/knowledge_base.py @@ -790,10 +790,16 @@ async def download_document( if not os.path.exists(document.file_path): raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="文件不存在") + import mimetypes + mime_type, _ = mimetypes.guess_type(document.original_filename or document.file_path) + if not mime_type: + mime_type = "application/octet-stream" + return FileResponse( document.file_path, filename=document.original_filename, - media_type="application/octet-stream" + media_type=mime_type, + content_disposition_type="inline" ) except HTTPException: diff --git a/backend/src/rag/document_loaders.py b/backend/src/rag/document_loaders.py index acfee97..3b5596d 100644 --- a/backend/src/rag/document_loaders.py +++ b/backend/src/rag/document_loaders.py @@ -1,5 +1,5 @@ """ -LangChain 1.0 文档加载器封装 + PDF图片提取 +LangChain 1.0 文档加载器封装 + PDF图片提取（Heron 版面检测 + 裁剪） """ import logging from typing import List, Optional, Dict @@ -19,11 +19,143 @@ logger = logging.getLogger(__name__) class PDFImageExtractor: - """使用pdf2image将每页PDF渲染为图片，用pymupdf提取页面文本""" + """使用 pdf2image 渲染页面 + Heron RT-DETR 检测 Picture 区域 + 裁剪""" + # 渲染参数 DEFAULT_DPI = 250 DEFAULT_TARGET_WIDTH = 2500 JPEG_QUALITY = 90 + CROP_PADDING = 10 + + # Heron 检测参数（参考 ZDTL config） + HERON_MODEL = "docling-project/docling-layout-heron" + PICTURE_THRESHOLD = 0.35 + NMS_IOU = 0.3 + MAX_AREA_RATIO = 0.45 + MIN_SIZE_W = 100 + MIN_SIZE_H = 80 + + @staticmethod + def _nms(boxes: list[dict], iou_threshold: float) -> list[dict]: + """IoU-based Non-Maximum Suppression""" + if not boxes: + return boxes + + import torch + + bboxes = torch.tensor([b["bbox"] for b in boxes], dtype=torch.float32) + scores = torch.tensor([b["score"] for b in boxes]) + + x1 = bboxes[:, 0] + y1 = bboxes[:, 1] + x2 = bboxes[:, 2] + y2 = bboxes[:, 3] + areas = (x2 - x1) * (y2 - y1) + + _, order = scores.sort(descending=True) + keep = [] + while order.numel() > 0: + if order.numel() == 1: + keep.append(order.item()) + break + i = order[0].item() + keep.append(i) + + xx1 = torch.max(x1[i], x1[order[1:]]) + yy1 = torch.max(y1[i], y1[order[1:]]) + xx2 = torch.min(x2[i], x2[order[1:]]) + yy2 = torch.min(y2[i], y2[order[1:]]) + + inter = (xx2 - xx1).clamp(min=0) * (yy2 - yy1).clamp(min=0) + union = areas[i] + areas[order[1:]] - inter + iou = inter / union + mask = iou <= iou_threshold + order = order[1:][mask] + + return [boxes[i] for i in keep] + + @staticmethod + def _detect_pictures( + pil_images: list, + model_name: str = None, + picture_threshold: float = None, + nms_iou: float = None, + max_area_ratio: float = None, + min_size_w: int = None, + min_size_h: int = None, + ) -> dict[int, list[dict]]: + """用 Heron RT-DETR 检测每页中的 Picture 区域 + + Returns: + {page_index: [{bbox: [x1,y1,x2,y2], score: float}, ...]} + """ + import torch + from transformers import RTDetrV2ForObjectDetection, RTDetrImageProcessor + + model_name = model_name or PDFImageExtractor.HERON_MODEL + picture_threshold = picture_threshold or PDFImageExtractor.PICTURE_THRESHOLD + nms_iou = nms_iou or PDFImageExtractor.NMS_IOU + max_area_ratio = max_area_ratio or PDFImageExtractor.MAX_AREA_RATIO + min_size_w = min_size_w or PDFImageExtractor.MIN_SIZE_W + min_size_h = min_size_h or PDFImageExtractor.MIN_SIZE_H + + logger.info(f"[ImageExtractor] 加载 Heron 模型: {model_name}") + processor = RTDetrImageProcessor.from_pretrained( + model_name, local_files_only=True + ) + model = RTDetrV2ForObjectDetection.from_pretrained( + model_name, local_files_only=True + ) + device = "cuda" if torch.cuda.is_available() else "cpu" + model.to(device) + model.eval() + logger.info(f"[ImageExtractor] 模型加载完成, device={device}") + + results: dict[int, list[dict]] = {} + + for idx, pil_img in enumerate(pil_images): + img = pil_img.convert("RGB") + W, H = img.size + inputs = processor(images=img, return_tensors="pt").to(device) + + with torch.no_grad(): + outputs = model(**inputs) + + detections = processor.post_process_object_detection( + outputs, threshold=0.1, target_sizes=[(H, W)] + )[0] + + pictures = [] + for score, label, box in zip( + detections["scores"], detections["labels"], detections["boxes"] + ): + s = score.item() + l = label.item() + x1, y1, x2, y2 = box.tolist() + area = (x2 - x1) * (y2 - y1) / (W * H) + w, h = x2 - x1, y2 - y1 + + if area > max_area_ratio or w < min_size_w or h < min_size_h: + continue + + cls = model.config.id2label.get(l, str(l)) + if cls == "picture" and s >= picture_threshold: + pictures.append({ + "bbox": [round(v, 1) for v in [x1, y1, x2, y2]], + "score": round(s, 3), + }) + + # NMS + 按位置排序 + pictures = PDFImageExtractor._nms(pictures, nms_iou) + pictures.sort(key=lambda b: (b["bbox"][1] // 200, b["bbox"][0])) + results[idx] = pictures + + logger.info( + f"[ImageExtractor] page {idx + 1}: " + f"{len(pictures)} pictures detected" + ) + + return results @staticmethod def extract_images( @@ -32,7 +164,9 @@ class PDFImageExtractor: dpi: int = None, target_width: int = None, ) -> List[dict]: - """将PDF每页渲染为一张图片 + """将PDF每页渲染为图片，用 Heron 检测并裁剪出图表区域 + + 若某页无检测到 Picture，则保存整页图作为兜底。 Args: file_path: PDF文件路径 @@ -46,10 +180,11 @@ class PDFImageExtractor: dpi = dpi or PDFImageExtractor.DEFAULT_DPI target_width = target_width or PDFImageExtractor.DEFAULT_TARGET_WIDTH quality = PDFImageExtractor.JPEG_QUALITY + padding = PDFImageExtractor.CROP_PADDING Path(output_dir).mkdir(parents=True, exist_ok=True) - # 用pymupdf提取每页文本（作为VLM上下文） + # Phase 1: pymupdf 提取页面文本 page_texts: list[str] = [] try: doc = fitz.open(file_path) @@ -59,58 +194,87 @@ class PDFImageExtractor: except Exception as e: logger.warning(f"[ImageExtractor] pymupdf文本提取失败: {e}") - # 用pdf2image渲染所有页面为图片 + # Phase 2: pdf2image 渲染页面 try: pil_images = convert_from_path(file_path, dpi=dpi, fmt="jpeg") except Exception as e: logger.error(f"[ImageExtractor] pdf2image渲染失败: {e}") return [] + # Phase 3: 缩放页面图 + scaled_images = [] + for pil_img in pil_images: + w, h = pil_img.size + if w > target_width: + ratio = target_width / w + new_h = int(h * ratio) + pil_img = pil_img.resize( + (target_width, new_h), Image.Resampling.LANCZOS + ) + scaled_images.append(pil_img) + + # Phase 4: Heron 检测 Picture 区域 + try: + detections = PDFImageExtractor._detect_pictures(scaled_images) + except Exception as e: + logger.error(f"[ImageExtractor] Heron检测失败: {e}") + return [] + + # Phase 5: 裁剪 + 保存 images = [] - for page_num, pil_img in enumerate(pil_images): - try: - w, h = pil_img.size - if w > target_width: - ratio = target_width / w - new_h = int(h * ratio) - pil_img = pil_img.resize( - (target_width, new_h), Image.Resampling.LANCZOS - ) + for page_idx, pil_img in enumerate(scaled_images): + page_num = page_idx + 1 + W, H = pil_img.size + context_text = "" + if page_idx < len(page_texts): + context_text = page_texts[page_idx][:600].strip() - filename = f"page{page_num + 1}.jpg" - output_path = Path(output_dir) / filename - pil_img.save(str(output_path), format="JPEG", quality=quality) + pics = detections.get(page_idx, []) - context_text = "" - if page_num < len(page_texts): - context_text = page_texts[page_num][:600].strip() + if pics: + # 裁剪检测到的 Picture 区域 + for fig_idx, pic in enumerate(pics): + try: + x1, y1, x2, y2 = pic["bbox"] + x1 = max(0, int(x1) - padding) + y1 = max(0, int(y1) - padding) + x2 = min(W, int(x2) + padding) + y2 = min(H, int(y2) + padding) - images.append({ - "path": str(output_path), - "filename": filename, - "page": page_num + 1, - "context_text": context_text, - "size": output_path.stat().st_size, - }) - logger.info( - f"[ImageExtractor] 页面渲染完成 page={page_num + 1} " - f"size={images[-1]['size']}" - ) - except Exception as e: - logger.warning( - f"[ImageExtractor] 页面处理失败 page={page_num + 1}: {e}" - ) + crop = pil_img.crop((x1, y1, x2, y2)) + filename = f"page{page_num}_fig{fig_idx + 1}.jpg" + output_path = Path(output_dir) / filename + crop.save(str(output_path), format="JPEG", quality=quality) + + images.append({ + "path": str(output_path), + "filename": filename, + "page": page_num, + "context_text": context_text, + "size": output_path.stat().st_size, + }) + logger.info( + f"[ImageExtractor] 裁剪 page={page_num} " + f"fig={fig_idx + 1} bbox={pic['bbox']} " + f"size={images[-1]['size']}" + ) + except Exception as e: + logger.warning( + f"[ImageExtractor] 裁剪失败 page={page_num} " + f"fig={fig_idx + 1}: {e}" + ) + else: continue logger.info( - f"[ImageExtractor] 共渲染 {len(images)} 页 from {file_path}" + f"[ImageExtractor] 完成: {len(images)} 张图片 from {file_path}" ) return images class DocumentLoaderFactory: """文档加载器工厂""" - + @staticmethod def get_loader(file_path: str, file_type: str): """根据文件类型获取对应的加载器""" @@ -120,21 +284,21 @@ class DocumentLoaderFactory: ".txt": TextLoader, ".md": UnstructuredMarkdownLoader, } - + loader_class = loaders.get(file_type) if not loader_class: raise ValueError(f"Unsupported file type: {file_type}") - + return loader_class(file_path) - + @staticmethod def load_document(file_path: str, file_type: str, metadata: Optional[dict] = None) -> List[Document]: """加载文档并添加元数据""" loader = DocumentLoaderFactory.get_loader(file_path, file_type) documents = loader.load() - + if metadata: for doc in documents: doc.metadata.update(metadata) - + return documents diff --git a/web/src/app/(main)/knowledge/[id]/page.tsx b/web/src/app/(main)/knowledge/[id]/page.tsx index ea7cd89..dfaa1d1 100644 --- a/web/src/app/(main)/knowledge/[id]/page.tsx +++ b/web/src/app/(main)/knowledge/[id]/page.tsx @@ -33,7 +33,7 @@ import { Settings } from "lucide-react"; import { formatFileSize, formatDate } from "@/lib/utils"; -import { knowledgeBaseAPI } from "@/lib/api"; +import { knowledgeBaseAPI, resolveImageUrl } from "@/lib/api"; import { KnowledgeBaseDetail, Document } from "@/types"; export default function KnowledgeBaseDetailPage() { @@ -158,13 +158,12 @@ export default function KnowledgeBaseDetailPage() { const handleViewDocument = async (doc: Document) => { try { const token = localStorage.getItem("auth_token"); - const res = await fetch(`/api/knowledge-bases/documents/${doc.id}/download`, { + const res = await fetch(resolveImageUrl(`/knowledge-bases/documents/${doc.id}/download`), { headers: { Authorization: `Bearer ${token}` }, }); - if (!res.ok) throw new Error("下载失败"); + if (!res.ok) throw new Error("获取文档失败"); const blob = await res.blob(); - // 从 Content-Disposition 提取文件名，或用 doc 信息拼接 const disposition = res.headers.get("Content-Disposition"); let filename = `${doc.title}${doc.file_type}`; if (disposition) { @@ -172,15 +171,11 @@ export default function KnowledgeBaseDetailPage() { if (match) filename = decodeURIComponent(match[1].replace(/["']/g, "")); } - // 触发浏览器下载 const url = URL.createObjectURL(blob); - const a = document.createElement("a"); - a.href = url; - a.download = filename; - a.click(); - URL.revokeObjectURL(url); + window.open(url, '_blank'); + setTimeout(() => URL.revokeObjectURL(url), 60000); } catch { - window.open(`/api/knowledge-bases/documents/${doc.id}/download`, '_blank'); + window.open(resolveImageUrl(`/knowledge-bases/documents/${doc.id}/download`), '_blank'); } }; diff --git a/web/src/components/chat/message-item.tsx b/web/src/components/chat/message-item.tsx index 20f90f6..7506880 100644 --- a/web/src/components/chat/message-item.tsx +++ b/web/src/components/chat/message-item.tsx @@ -113,7 +113,18 @@ export default function MessageItem({ message, selectedModel }: MessageItemProps const handleCopy = async () => { try { - await navigator.clipboard.writeText(message.content); + if (navigator.clipboard && window.isSecureContext) { + await navigator.clipboard.writeText(message.content); + } else { + const textarea = document.createElement("textarea"); + textarea.value = message.content; + textarea.style.position = "fixed"; + textarea.style.left = "-9999px"; + document.body.appendChild(textarea); + textarea.select(); + document.execCommand("copy"); + document.body.removeChild(textarea); + } toast.success("已复制"); } catch { toast.error("复制失败"); diff --git a/web/src/store/chat.ts b/web/src/store/chat.ts index 4a6d709..3e90a53 100644 --- a/web/src/store/chat.ts +++ b/web/src/store/chat.ts @@ -339,6 +339,15 @@ export const useChatStore = create((set, get) => ({ scheduleFlush(); }, (_sessionId: number, messageId?: number, userMessageId?: number) => { + // 刷出缓冲区中剩余的内容 + if (rafId !== null) { + cancelAnimationFrame(rafId); + rafId = null; + } + if (contentBuffer) { + flushBuffer(); + } + if (messageId) { set((state) => ({ messages: state.messages.map(msg => {