Compare commits

...

2 Commits

Author SHA1 Message Date
pengxiao 6395ee3b49 feat: replace pymupdf image extraction with pdf2image + Heron layout detection
PDF image extraction now renders pages with pdf2image, detects figure
regions using docling-layout-heron (RT-DETRv2), and crops only the
detected pictures. Removes full-page fallback for text-only pages.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-06-02 22:38:31 +08:00
pengxiao bbf6b921af fix: adapt frontend image paths for backend port 8002 and use blob download
Backend port changed from 8000 to 8002 due to port conflict (frpc on Windows).
- Add resolveImageUrl() to convert relative image URLs to full backend URLs
- Fix all spatial page image src and download handlers to use backend address
- Fix chat message image rendering to use correct port
- Switch download to blob-based approach to support cross-origin file saving

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-06-02 20:52:48 +08:00
14 changed files with 514 additions and 144 deletions
+3
View File
@@ -71,3 +71,6 @@ Desktop.ini
htmlcov/
officefile
data/1法律
data/data_backup
+1
View File
@@ -43,6 +43,7 @@ dependencies = [
"docx2txt>=0.9",
"pypdf>=6.12.0",
"pymupdf>=1.27.2.3",
"pdf2image>=1.17.0",
"opencv-python-headless>=4.13.0.92",
]
+140
View File
@@ -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):
"""阶段1Heron 版面检测 + 裁剪"""
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):
"""阶段2VLM图片描述"""
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()
+7 -1
View File
@@ -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:
+4
View File
@@ -128,6 +128,10 @@ class Settings(BaseSettings):
# 全局配置实例
settings = Settings()
# 将 HF_ENDPOINT 写入环境变量,供 sentence_transformers/huggingface_hub 使用
if settings.hf_endpoint and not os.environ.get("HF_ENDPOINT"):
os.environ["HF_ENDPOINT"] = settings.hf_endpoint
def get_settings() -> Settings:
"""获取配置实例"""
+245 -89
View File
@@ -1,12 +1,12 @@
"""
LangChain 1.0 文档加载器封装 + PDF图片提取
LangChain 1.0 文档加载器封装 + PDF图片提取Heron 版面检测 + 裁剪)
"""
import logging
import io
from typing import List, Optional, Dict
from pathlib import Path
import fitz # pymupdf
import fitz # pymupdf — 仅用于提取页面文本
from PIL import Image
from pdf2image import convert_from_path
from langchain_community.document_loaders import (
PyPDFLoader,
Docx2txtLoader,
@@ -19,106 +19,262 @@ logger = logging.getLogger(__name__)
class PDFImageExtractor:
"""使用pymupdf从PDF中提取内嵌图片"""
"""使用 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 extract_images(file_path: str, output_dir: str) -> List[dict]:
"""提取PDF中所有图片,返回图片元数据列表
def _nms(boxes: list[dict], iou_threshold: float) -> list[dict]:
"""IoU-based Non-Maximum Suppression"""
if not boxes:
return boxes
对正常页面提取内嵌图片,对瓦片式页面(>50个小图片)渲染整页截图。
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}, ...]}
"""
Path(output_dir).mkdir(parents=True, exist_ok=True)
images = []
doc = fitz.open(file_path)
import torch
from transformers import RTDetrV2ForObjectDetection, RTDetrImageProcessor
TILE_THRESHOLD = 50
vlm_formats = {"png", "jpg", "jpeg", "webp", "gif"}
dpi = 150 # 页面渲染 DPI
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
for page_num in range(len(doc)):
page = doc[page_num]
page_text = page.get_text("text")
image_list = page.get_images(full=True)
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}")
if len(image_list) > TILE_THRESHOLD:
# 瓦片式页面:渲染整页为一张完整图片
try:
pix = page.get_pixmap(dpi=dpi)
img_bytes = pix.tobytes("png")
filename = f"page{page_num+1}_full.png"
output_path = Path(output_dir) / filename
output_path.write_bytes(img_bytes)
results: dict[int, list[dict]] = {}
images.append({
"path": str(output_path),
"filename": filename,
"page": page_num + 1,
"context_text": page_text[:600].strip(),
"size": len(img_bytes),
})
logger.info(f"[ImageExtractor] 瓦片页面渲染为整图 page={page_num+1} size={len(img_bytes)}")
except Exception as e:
logger.warning(f"[ImageExtractor] 页面渲染失败 page={page_num+1}: {e}")
continue
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)
# 正常页面:提取内嵌图片
for img_idx, img_info in enumerate(image_list):
xref = img_info[0]
try:
base_image = doc.extract_image(xref)
image_bytes = base_image["image"]
ext = base_image["ext"]
with torch.no_grad():
outputs = model(**inputs)
# 过滤太小的碎片/图标/装饰
if len(image_bytes) < 2048:
continue
w = base_image.get("width", 0)
h = base_image.get("height", 0)
if w < 150 or h < 150:
continue
detections = processor.post_process_object_detection(
outputs, threshold=0.1, target_sizes=[(H, W)]
)[0]
# 图片周围文本
context_start = max(0, page_text.find(
page_text[:len(page_text)//2]) if len(page_text) > 600
else 0)
context_text = page_text[context_start:context_start+600].strip()
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
# 不支持的格式(jpx/jpeg2000等)转为 PNG
if ext.lower() not in vlm_formats:
try:
img = Image.open(io.BytesIO(image_bytes))
if img.mode in ("CMYK", "P"):
img = img.convert("RGB")
buf = io.BytesIO()
img.save(buf, format="PNG")
image_bytes = buf.getvalue()
ext = "png"
except Exception as e:
logger.warning(f"[ImageExtractor] 格式转换失败 page={page_num+1} img={img_idx}: {e}")
continue
filename = f"page{page_num+1}_img{img_idx+1}.{ext}"
output_path = Path(output_dir) / filename
output_path.write_bytes(image_bytes)
images.append({
"path": str(output_path),
"filename": filename,
"page": page_num + 1,
"context_text": context_text,
"size": len(image_bytes),
})
except Exception as e:
logger.warning(f"[ImageExtractor] 提取图片失败 page={page_num+1} img={img_idx}: {e}")
if area > max_area_ratio or w < min_size_w or h < min_size_h:
continue
doc.close()
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(
file_path: str,
output_dir: str,
dpi: int = None,
target_width: int = None,
) -> List[dict]:
"""将PDF每页渲染为图片,用 Heron 检测并裁剪出图表区域
若某页无检测到 Picture,则保存整页图作为兜底。
Args:
file_path: PDF文件路径
output_dir: 图片输出目录
dpi: 渲染DPI(默认250
target_width: 缩放目标宽度像素(默认2500)
Returns:
图片元数据列表,每个dict包含 path, filename, page, context_text, size
"""
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)
# Phase 1: pymupdf 提取页面文本
page_texts: list[str] = []
try:
doc = fitz.open(file_path)
for page_num in range(len(doc)):
page_texts.append(doc[page_num].get_text("text"))
doc.close()
except Exception as e:
logger.warning(f"[ImageExtractor] pymupdf文本提取失败: {e}")
# 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_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()
pics = detections.get(page_idx, [])
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)
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}"
)
return images
class DocumentLoaderFactory:
"""文档加载器工厂"""
@staticmethod
def get_loader(file_path: str, file_type: str):
"""根据文件类型获取对应的加载器"""
@@ -128,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
+15 -1
View File
@@ -1,5 +1,5 @@
version = 1
revision = 2
revision = 3
requires-python = ">=3.12"
resolution-markers = [
"python_full_version >= '3.13'",
@@ -441,6 +441,7 @@ dependencies = [
{ name = "opencv-python-headless" },
{ name = "pandas" },
{ name = "passlib", extra = ["bcrypt"] },
{ name = "pdf2image" },
{ name = "psycopg", extra = ["binary"] },
{ name = "pydantic", extra = ["email"] },
{ name = "pydantic-settings" },
@@ -492,6 +493,7 @@ requires-dist = [
{ name = "opencv-python-headless", specifier = ">=4.13.0.92" },
{ name = "pandas", specifier = ">=2.2.3" },
{ name = "passlib", extras = ["bcrypt"], specifier = ">=1.7.4" },
{ name = "pdf2image", specifier = ">=1.17.0" },
{ name = "psycopg", extras = ["binary"], specifier = ">=3.1.0" },
{ name = "pydantic", extras = ["email"], specifier = ">=2.5.0" },
{ name = "pydantic-settings", specifier = ">=2.1.0" },
@@ -2396,6 +2398,18 @@ wheels = [
{ url = "https://mirrors.aliyun.com/pypi/packages/cc/20/ff623b09d963f88bfde16306a54e12ee5ea43e9b597108672ff3a408aad6/pathspec-0.12.1-py3-none-any.whl", hash = "sha256:a0d503e138a4c123b27490a4f7beda6a01c6f288df0e4a8b79c7eb0dc7b4cc08" },
]
[[package]]
name = "pdf2image"
version = "1.17.0"
source = { registry = "https://mirrors.aliyun.com/pypi/simple/" }
dependencies = [
{ name = "pillow" },
]
sdist = { url = "https://mirrors.aliyun.com/pypi/packages/00/d8/b280f01045555dc257b8153c00dee3bc75830f91a744cd5f84ef3a0a64b1/pdf2image-1.17.0.tar.gz", hash = "sha256:eaa959bc116b420dd7ec415fcae49b98100dda3dd18cd2fdfa86d09f112f6d57" }
wheels = [
{ url = "https://mirrors.aliyun.com/pypi/packages/62/33/61766ae033518957f877ab246f87ca30a85b778ebaad65b7f74fa7e52988/pdf2image-1.17.0-py3-none-any.whl", hash = "sha256:ecdd58d7afb810dffe21ef2b1bbc057ef434dabbac6c33778a38a3f7744a27e2" },
]
[[package]]
name = "pillow"
version = "11.3.0"
+6 -11
View File
@@ -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');
}
};
@@ -12,7 +12,7 @@ import { Badge } from "@/components/ui/badge";
import { Separator } from "@/components/ui/separator";
import { Slider } from "@/components/ui/slider";
import { Loader2, Download, Copy, Upload, Edit3, Wand2, Expand, Palette, Image as ImageIcon, X } from "lucide-react";
import { imageAPI } from "@/lib/api";
import { imageAPI, resolveImageUrl } from "@/lib/api";
import { toast } from "sonner";
interface EditResult {
@@ -312,14 +312,22 @@ export default function ImageToImagePage() {
}
};
const handleDownload = (imageUrl: string, imageId: string) => {
const link = document.createElement("a");
link.href = imageUrl;
link.download = `edited-image-${imageId}.png`;
document.body.appendChild(link);
link.click();
document.body.removeChild(link);
toast.success("图像下载成功");
const handleDownload = async (imageUrl: string, imageId: string) => {
try {
const res = await fetch(imageUrl);
const blob = await res.blob();
const url = URL.createObjectURL(blob);
const link = document.createElement("a");
link.href = url;
link.download = `edited-image-${imageId}.png`;
document.body.appendChild(link);
link.click();
document.body.removeChild(link);
URL.revokeObjectURL(url);
toast.success("图像下载成功");
} catch {
toast.error("图像下载失败");
}
};
return (
@@ -595,7 +603,7 @@ export default function ImageToImagePage() {
<div className="space-y-4">
<div className="relative group">
<img
src={editResult.url}
src={resolveImageUrl(editResult.url)}
alt="Edited image"
className="w-full h-64 object-cover rounded-lg"
/>
@@ -604,7 +612,7 @@ export default function ImageToImagePage() {
<Button
size="sm"
variant="secondary"
onClick={() => handleDownload(editResult.url, editResult.id)}
onClick={() => handleDownload(resolveImageUrl(editResult.url), editResult.id)}
>
<Download className="h-3 w-3" />
</Button>
@@ -648,7 +656,7 @@ export default function ImageToImagePage() {
>
<div className="relative">
<img
src={variation.url}
src={resolveImageUrl(variation.url)}
alt={`Variation ${index + 1}`}
className="w-full h-48 object-cover rounded-lg"
/>
@@ -657,7 +665,7 @@ export default function ImageToImagePage() {
<Button
size="sm"
variant="secondary"
onClick={() => handleDownload(variation.url, variation.id)}
onClick={() => handleDownload(resolveImageUrl(variation.url), variation.id)}
>
<Download className="h-3 w-3" />
</Button>
+23 -15
View File
@@ -32,7 +32,7 @@ import {
X
} from "lucide-react";
import LoadingSpinner from "@/components/ui/loading-spinner";
import { imageAPI } from "@/lib/api";
import { imageAPI, resolveImageUrl } from "@/lib/api";
import { toast } from "sonner";
import { motion } from "framer-motion";
@@ -350,14 +350,22 @@ export default function SpatialPage() {
}
};
const handleDownload = (imageUrl: string, imageId: string) => {
const link = document.createElement("a");
link.href = imageUrl;
link.download = `image-${imageId}.png`;
document.body.appendChild(link);
link.click();
document.body.removeChild(link);
toast.success("图像下载成功");
const handleDownload = async (imageUrl: string, imageId: string) => {
try {
const res = await fetch(imageUrl);
const blob = await res.blob();
const url = URL.createObjectURL(blob);
const link = document.createElement("a");
link.href = url;
link.download = `image-${imageId}.png`;
document.body.appendChild(link);
link.click();
document.body.removeChild(link);
URL.revokeObjectURL(url);
toast.success("图像下载成功");
} catch {
toast.error("图像下载失败");
}
};
if (authLoading) {
@@ -580,7 +588,7 @@ export default function SpatialPage() {
<Card className="overflow-hidden">
<div className="aspect-square relative bg-muted">
<img
src={image.url}
src={resolveImageUrl(image.url)}
alt={`Generated image ${index + 1}`}
className="w-full h-full object-cover"
loading="lazy"
@@ -599,7 +607,7 @@ export default function SpatialPage() {
variant="outline"
size="sm"
className="w-full"
onClick={() => handleDownload(image.url, image.id)}
onClick={() => handleDownload(resolveImageUrl(image.url), image.id)}
>
<Download className="h-4 w-4 mr-2" />
@@ -852,7 +860,7 @@ export default function SpatialPage() {
<div className="space-y-4">
<div className="relative group bg-muted rounded-lg overflow-hidden">
<img
src={editResult.url}
src={resolveImageUrl(editResult.url)}
alt="Edited image"
className="w-full h-64 object-contain"
loading="lazy"
@@ -871,7 +879,7 @@ export default function SpatialPage() {
variant="outline"
size="sm"
className="w-full"
onClick={() => handleDownload(editResult.url, editResult.id)}
onClick={() => handleDownload(resolveImageUrl(editResult.url), editResult.id)}
>
<Download className="h-4 w-4 mr-2" />
@@ -907,7 +915,7 @@ export default function SpatialPage() {
<Card className="overflow-hidden">
<div className="relative bg-muted">
<img
src={variation.url}
src={resolveImageUrl(variation.url)}
alt={`Variation ${index + 1}`}
className="w-full h-48 object-cover"
loading="lazy"
@@ -922,7 +930,7 @@ export default function SpatialPage() {
variant="outline"
size="sm"
className="w-full"
onClick={() => handleDownload(variation.url, variation.id)}
onClick={() => handleDownload(resolveImageUrl(variation.url), variation.id)}
>
<Download className="h-4 w-4 mr-2" />
下载图片
@@ -11,7 +11,7 @@ import { Textarea } from "@/components/ui/textarea";
import { Badge } from "@/components/ui/badge";
import { Separator } from "@/components/ui/separator";
import { Loader2, Download, Copy, RefreshCw, Sparkles, Image as ImageIcon } from "lucide-react";
import { imageAPI } from "@/lib/api";
import { imageAPI, resolveImageUrl } from "@/lib/api";
import { toast } from "sonner";
interface GeneratedImage {
@@ -95,14 +95,22 @@ export default function TextToImagePage() {
}
};
const handleDownload = (imageUrl: string, imageId: string) => {
const link = document.createElement("a");
link.href = imageUrl;
link.download = `generated-image-${imageId}.png`;
document.body.appendChild(link);
link.click();
document.body.removeChild(link);
toast.success("图像下载成功");
const handleDownload = async (imageUrl: string, imageId: string) => {
try {
const res = await fetch(imageUrl);
const blob = await res.blob();
const url = URL.createObjectURL(blob);
const link = document.createElement("a");
link.href = url;
link.download = `generated-image-${imageId}.png`;
document.body.appendChild(link);
link.click();
document.body.removeChild(link);
URL.revokeObjectURL(url);
toast.success("图像下载成功");
} catch {
toast.error("图像下载失败");
}
};
const handleCopyPrompt = (prompt: string) => {
@@ -316,7 +324,7 @@ export default function TextToImagePage() {
<Card className="overflow-hidden">
<div className="aspect-square relative">
<img
src={image.url}
src={resolveImageUrl(image.url)}
alt={`Generated image ${index + 1}`}
className="w-full h-full object-cover"
/>
@@ -326,7 +334,7 @@ export default function TextToImagePage() {
<Button
size="sm"
variant="secondary"
onClick={() => handleDownload(image.url, image.id)}
onClick={() => handleDownload(resolveImageUrl(image.url), image.id)}
>
<Download className="h-3 w-3" />
</Button>
+13 -2
View File
@@ -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("复制失败");
@@ -254,7 +265,7 @@ export default function MessageItem({ message, selectedModel }: MessageItemProps
),
img: ({ src, alt }: any) => {
const resolvedSrc = src && src.startsWith("/")
? `${process.env.NEXT_PUBLIC_API_URL || `${window.location.protocol}//${window.location.hostname}:8000`}${src}`
? `${process.env.NEXT_PUBLIC_API_URL || `${window.location.protocol}//${window.location.hostname}:8002`}${src}`
: src;
return (
<a href={resolvedSrc} target="_blank" rel="noopener noreferrer" className="block my-3 group">
+8 -1
View File
@@ -12,7 +12,14 @@ import type {
} from "@/types";
// API基础配置
const API_BASE_URL = process.env.NEXT_PUBLIC_API_URL || (typeof window !== 'undefined' ? `${window.location.protocol}//${window.location.hostname}:8000` : "http://127.0.0.1:8000");
const API_BASE_URL = process.env.NEXT_PUBLIC_API_URL || (typeof window !== 'undefined' ? `${window.location.protocol}//${window.location.hostname}:8002` : "http://127.0.0.1:8002");
// 将后端返回的相对路径图片URL转为完整的后端地址
export function resolveImageUrl(url: string): string {
if (!url) return url;
if (url.startsWith("http://") || url.startsWith("https://") || url.startsWith("data:")) return url;
return `${API_BASE_URL}${url.startsWith("/") ? "" : "/"}${url}`;
}
// 请求拦截器
async function apiRequest<T>(
+9
View File
@@ -339,6 +339,15 @@ export const useChatStore = create<ChatStore>((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 => {