""" PDF图片提取器 - 从PDF中提取嵌入的图片并与题目关联 专门处理数学试卷中的几何图形、函数图像、表格等 """ import os import time import hashlib import tempfile from typing import List, Dict, Tuple, Optional import fitz # PyMuPDF from PIL import Image import uuid class PDFImageExtractor: def __init__(self, output_dir: str = "uploads/images"): self.output_dir = output_dir # 确保输出目录存在 os.makedirs(output_dir, exist_ok=True) # 图片提取配置 self.min_image_size = (50, 50) # 最小图片尺寸，过滤掉太小的图片 self.max_image_size = (2000, 2000) # 最大图片尺寸 self.supported_formats = ['png', 'jpg', 'jpeg', 'gif', 'bmp'] def extract_images_from_pdf(self, pdf_path: str) -> List[Dict]: """ 从PDF中提取所有图片 返回图片信息列表，包含图片路径、位置、尺寸等 """ print(f"🖼️ 开始从PDF提取图片: {pdf_path}") start_time = time.time() try: doc = fitz.open(pdf_path) extracted_images = [] for page_num in range(len(doc)): page = doc[page_num] print(f"📄 处理第{page_num + 1}页...") # 获取页面中的图片 image_list = page.get_images(full=True) for img_index, img in enumerate(image_list): # 提取图片信息 xref = img[0] # 图片引用 # 获取图片数据 base_image = doc.extract_image(xref) image_bytes = base_image["image"] image_ext = base_image["ext"] # 验证图片格式 if image_ext.lower() not in self.supported_formats: continue # 生成唯一文件名 image_hash = hashlib.md5(image_bytes).hexdigest() filename = f"page_{page_num + 1}_img_{img_index + 1}_{image_hash[:8]}.{image_ext}" image_path = os.path.join(self.output_dir, filename) # 保存图片 with open(image_path, "wb") as img_file: img_file.write(image_bytes) # 获取图片在页面中的位置 image_rects = page.get_image_rects(xref) # 验证图片尺寸 try: with Image.open(image_path) as pil_img: width, height = pil_img.size # 过滤太小或太大的图片 if (width < self.min_image_size[0] or height < self.min_image_size[1] or width > self.max_image_size[0] or height > self.max_image_size[1]): os.remove(image_path) # 删除不符合条件的图片 continue except Exception as e: print(f"⚠️ 图片验证失败: {e}") if os.path.exists(image_path): os.remove(image_path) continue # 记录图片信息 image_info = { "filename": filename, "path": image_path, "relative_path": f"images/{filename}", # 用于前端访问 "page": page_num + 1, "width": width, "height": height, "format": image_ext.upper(), "size_bytes": len(image_bytes), "hash": image_hash, "rects": [] # 图片在页面中的位置 } # 添加位置信息 for rect in image_rects: image_info["rects"].append({ "x0": rect.x0, "y0": rect.y0, "x1": rect.x1, "y1": rect.y1, "width": rect.width, "height": rect.height }) extracted_images.append(image_info) print(f"✅ 提取图片: {filename} ({width}x{height})") doc.close() total_time = time.time() - start_time print(f"🎉 图片提取完成: {len(extracted_images)}张图片, 耗时{total_time:.2f}秒") return extracted_images except Exception as e: print(f"❌ PDF图片提取失败: {str(e)}") return [] def associate_images_with_questions(self, images: List[Dict], questions: List[Dict]) -> List[Dict]: """ 将图片与题目关联 基于页面位置和题目边界框进行智能匹配 """ print(f"🔗 开始关联图片与题目: {len(images)}张图片, {len(questions)}道题目") # 用于跟踪已分配的图片 used_images = set() # 按页码组织题目 questions_by_page = {} for i, question in enumerate(questions): page = question.get("page", 1) if page not in questions_by_page: questions_by_page[page] = [] questions_by_page[page].append((i, question)) # 为每页进行图片关联 for page_num, page_questions in questions_by_page.items(): # 获取当前页的图片 page_images = [img for img in images if img["page"] == page_num] if not page_images: continue # 按图片在页面中的垂直位置排序（从上到下） page_images_sorted = sorted(page_images, key=lambda x: min(rect["y0"] for rect in x["rects"]) if x["rects"] else 0) # 为每个题目进行图片分配 for question_idx, question in page_questions: question_images = [] # 获取题目的边界框信息 question_bbox = self._get_question_bbox(question) # 遍历页面中的图片，寻找合适的匹配 for img_idx, image in enumerate(page_images_sorted): img_filename = image["filename"] # 跳过已分配的图片 if img_filename in used_images: continue # 检查图片是否适合当前题目 if self._should_associate_image(question, image, question_bbox): question_images.append({ "filename": image["filename"], "path": image["relative_path"], "width": image["width"], "height": image["height"], "format": image["format"] }) used_images.add(img_filename) print(f"📝 题目{question.get('number', '?')}关联图片: {img_filename}") break # 每个题目最多关联一张图片，除非明确需要多张 question["images"] = question_images if not question_images: print(f"📝 题目{question.get('number', '?')}未关联图片") return questions def _get_question_bbox(self, question: Dict) -> Dict: """ 获取题目的边界框信息 """ # 尝试从不同字段获取边界框 bbox = question.get("bbox") if bbox and isinstance(bbox, dict): return bbox # 尝试从page_spans获取 page_spans = question.get("page_spans") if page_spans and len(page_spans) > 0: span = page_spans[0] if isinstance(span, dict) and "bbox" in span: return span["bbox"] # 默认边界框（整个页面） return {"x0": 0, "y0": 0, "x1": 1000, "y1": 1000} def _should_associate_image(self, question: Dict, image: Dict, question_bbox: Dict) -> bool: """ 判断图片是否应该关联到当前题目 """ # 策略1: 基于几何位置匹配 if image.get("rects"): for rect in image["rects"]: # 检查图片是否在题目的合理范围内 # 考虑页面布局和图片通常的放置位置 img_center_y = (rect["y0"] + rect["y1"]) / 2 question_center_y = (question_bbox.get("y0", 0) + question_bbox.get("y1", 1000)) / 2 # 如果图片中心在题目中心的合理范围内（上下200点） if abs(img_center_y - question_center_y) < 200: return True # 如果图片在题目下方（通常图片紧跟题目） if rect["y0"] > question_bbox.get("y1", 0) and rect["y0"] - question_bbox.get("y1", 0) < 300: return True # 策略2: 基于题目内容判断 question_text = question.get("text", "").lower() question_type = question.get("type", "") # 选择题通常较少有图片 if question_type in ["single_choice", "multiple_choice"]: # 只有明确提到图形时才关联图片 if any(keyword in question_text for keyword in ["如图", "下图", "图形", "图像", "坐标", "函数图像"]): return True # 解答题、证明题等更可能有图片 elif question_type in ["essay", "calculation", "proof"]: return True # 默认不关联（更保守的策略） return False def extract_and_associate(self, pdf_path: str, questions: List[Dict]) -> Tuple[List[Dict], List[Dict]]: """ 提取PDF图片并与题目关联 返回 (更新后的题目列表, 图片列表) """ # 1. 提取图片 images = self.extract_images_from_pdf(pdf_path) # 2. 关联图片与题目 updated_questions = self.associate_images_with_questions(images, questions) return updated_questions, images def clean_unused_images(self, used_images: List[str]): """ 清理未使用的图片文件 """ try: if not os.path.exists(self.output_dir): return # 获取所有图片文件 all_files = os.listdir(self.output_dir) # 删除未使用的文件 for filename in all_files: if filename not in used_images: file_path = os.path.join(self.output_dir, filename) try: os.remove(file_path) print(f"🗑️ 清理未使用图片: {filename}") except Exception as e: print(f"⚠️ 清理文件失败: {filename}, {e}") except Exception as e: print(f"⚠️ 清理图片失败: {e}") # 测试函数 def test_image_extractor(): """ 测试图片提取器 """ pdf_path = "/Users/brunogao/Downloads/深圳市八年级（上）期中数学试卷1.pdf" if not os.path.exists(pdf_path): print(f"❌ 测试文件不存在: {pdf_path}") return extractor = PDFImageExtractor() # 模拟一些题目数据 sample_questions = [ {"number": "1", "page": 1, "text": "题目1内容"}, {"number": "2", "page": 1, "text": "题目2内容"}, {"number": "3", "page": 2, "text": "题目3内容"}, ] # 提取并关联图片 updated_questions, images = extractor.extract_and_associate(pdf_path, sample_questions) print(f"\n📊 测试结果:") print(f"- 提取图片数: {len(images)}") print(f"- 处理题目数: {len(updated_questions)}") for question in updated_questions: if question.get("images"): print(f"题目{question['number']}: {len(question['images'])}张图片") for img in question["images"]: print(f" - {img['filename']} ({img['width']}x{img['height']})") if __name__ == "__main__": test_image_extractor()