content: # !pip install PyMuPDF
  # !pip install matplotlib
  
  import os
  import fitz
  from PIL import Image
  
  def pdf2images(pdf_file):
      '''将 PDF 每页转成一个 PNG 图像'''
      # 保存路径为原 PDF 文件名（不含扩展名）
      output_directory_path, _ = os.path.splitext(pdf_file)
      
      if not os.path.exists(output_directory_path):
          os.makedirs(output_directory_path)
      
      # 加载 PDF 文件
      pdf_document = fitz.open(pdf_file)
      
      # 每页转一张图
      for page_number in range(pdf_document.page_count):
          # 取一页
          page = pdf_document[page_number]
      
          # 转图像
          pix = page.get_pixmap()
      
          # 从位图创建 PNG 对象
          image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
      
          # 保存 PNG 文件
          image.save(f"./{output_directory_path}/page_{page_number + 1}.png")
      
      # 关闭 PDF 文件
      pdf_document.close()
      
  from PIL import Image
  import os
  import matplotlib.pyplot as plt
  
  def show_images(dir_path):
      '''显示目录下的 PNG 图像'''
      for file in os.listdir(dir_path):
          if file.endswith('.png'):
              # 打开图像
              img = Image.open(os.path.join(dir_path, file)) 
  
              # 显示图像
              plt.imshow(img)
              plt.axis('off')  # 不显示坐标轴
              plt.show()
              
  pdf2images("llama2_page8.pdf")
  show_images("llama2_page8")
  
  # 识别文档中的表格
  
  class MaxResize(object):
      '''缩放图像'''
      def __init__(self, max_size=800):
          self.max_size = max_size
  
      def __call__(self, image):
          width, height = image.size
          current_max_size = max(width, height)
          scale = self.max_size / current_max_size
          resized_image = image.resize(
              (int(round(scale * width)), int(round(scale * height)))
          )
  
          return resized_image
  
  import torchvision.transforms as transforms
  
  # 图像预处理
  detection_transform = transforms.Compose(
      [
          MaxResize(800),
          transforms.ToTensor(),
          transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
      ]
  )
  
  from transformers import AutoModelForObjectDetection
  
  # 加载 TableTransformer 模型
  model = AutoModelForObjectDetection.from_pretrained(
      "microsoft/table-transformer-detection"
  )
  
  # 识别后的坐标换算与后处理
  
  def box_cxcywh_to_xyxy(x):
      '''坐标转换'''
      x_c, y_c, w, h = x.unbind(-1)
      b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]
      return torch.stack(b, dim=1)
  
  
  def rescale_bboxes(out_bbox, size):
      '''区域缩放'''
      width, height = size
      boxes = box_cxcywh_to_xyxy(out_bbox)
      boxes = boxes * torch.tensor(
          [width, height, width, height], dtype=torch.float32
      )
      return boxes
  
  
  def outputs_to_objects(outputs, img_size, id2label):
      '''从模型输出中取定位框坐标'''
      m = outputs.logits.softmax(-1).max(-1)
      pred_labels = list(m.indices.detach().cpu().numpy())[0]
      pred_scores = list(m.values.detach().cpu().numpy())[0]
      pred_bboxes = outputs["pred_boxes"].detach().cpu()[0]
      pred_bboxes = [
          elem.tolist() for elem in rescale_bboxes(pred_bboxes, img_size)
      ]
  
      objects = []
      for label, score, bbox in zip(pred_labels, pred_scores, pred_bboxes):
          class_label = id2label[int(label)]
          if not class_label == "no object":
              objects.append(
                  {
                      "label": class_label,
                      "score": float(score),
                      "bbox": [float(elem) for elem in bbox],
                  }
              )
  
      return objects
    
  import torch
  
  # 识别表格，并将表格部分单独存为图像文件
  
  def detect_and_crop_save_table(file_path):
      # 加载图像（PDF页）    
      image = Image.open(file_path)
  
      filename, _ = os.path.splitext(os.path.basename(file_path))
  
      # 输出路径
      cropped_table_directory = os.path.join(os.path.dirname(file_path), "table_images")
  
      if not os.path.exists(cropped_table_directory):
          os.makedirs(cropped_table_directory)
  
      # 预处理
      pixel_values = detection_transform(image).unsqueeze(0)
  
      # 识别表格
      with torch.no_grad():
          outputs = model(pixel_values)
  
      # 后处理，得到表格子区域
      id2label = model.config.id2label
      id2label[len(model.config.id2label)] = "no object"
      detected_tables = outputs_to_objects(outputs, image.size, id2label)
  
      print(f"number of tables detected {len(detected_tables)}")
  
      for idx in range(len(detected_tables)):
          # 将识别从的表格区域单独存为图像
          cropped_table = image.crop(detected_tables[idx]["bbox"])
          cropped_table.save(os.path.join(cropped_table_directory,f"{filename}_{idx}.png"))
          
  detect_and_crop_save_table("llama2_page8/page_1.png")
  show_images("llama2_page8/table_images")
  
  
  # 基于 GPT-4 Vision API 做表格问答
  
  import base64
  from openai import OpenAI
  
  client = OpenAI()
  
  def encode_image(image_path):
    with open(image_path, "rb") as image_file:
      return base64.b64encode(image_file.read()).decode('utf-8')
  
  def image_qa(query, image_path):
      base64_image = encode_image(image_path)
      response = client.chat.completions.create(
          model="gpt-4o",
          temperature=0,
          seed=42,
          messages=[{
              "role": "user",
                "content": [
                    {"type": "text", "text": query},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}",
                        },
                    },
                ],
          }],
      )
  
      return response.choices[0].message.content
    
    
  response = image_qa("哪个模型在AGI Eval数据集上表现最好。得分多少","llama2_page8/table_images/page_1_0.png")
  print(response)
  
  
  # 用 GPT-4 Vision 生成表格（图像）描述，并向量化用于检索
  
  
  import chromadb
  from chromadb.config import Settings
  
  
  class NewVectorDBConnector:
      def __init__(self, collection_name, embedding_fn):
          chroma_client = chromadb.Client(Settings(allow_reset=True))
  
          # 为了演示，实际不需要每次 reset()
          chroma_client.reset()
  
          # 创建一个 collection
          self.collection = chroma_client.get_or_create_collection(
              name=collection_name)
          self.embedding_fn = embedding_fn
  
      def add_documents(self, documents):
          '''向 collection 中添加文档与向量'''
          self.collection.add(
              embeddings=self.embedding_fn(documents),  # 每个文档的向量
              documents=documents,  # 文档的原文
              ids=[f"id{i}" for i in range(len(documents))]  # 每个文档的 id
          )
  
      def add_images(self, image_paths):
          '''向 collection 中添加图像'''
          documents = [
              image_qa("请简要描述图片中的信息",image)
              for image in image_paths
          ]
          self.collection.add(
              embeddings=self.embedding_fn(documents),  # 每个文档的向量
              documents=documents,  # 文档的原文
              ids=[f"id{i}" for i in range(len(documents))],  # 每个文档的 id
              metadatas=[{"image": image} for image in image_paths] # 用 metadata 标记源图像路径
          )
  
      def search(self, query, top_n):
          '''检索向量数据库'''
          results = self.collection.query(
              query_embeddings=self.embedding_fn([query]),
              n_results=top_n
          )
          return results
        
        
  images = []
  dir_path = "llama2_page8/table_images"
  for file in os.listdir(dir_path):
      if file.endswith('.png'):
          # 打开图像
          images.append(os.path.join(dir_path, file))
  
  new_db_connector = NewVectorDBConnector("table_demo",get_embeddings)
  new_db_connector.add_images(images)
  
  query  = "哪个模型在AGI Eval数据集上表现最好。得分多少"
  
  results = new_db_connector.search(query, 1)
  metadata = results["metadatas"][0]
  print("====检索结果====")
  print(metadata)
  print("====回复====")
  response = image_qa(query,metadata[0]["image"])
  print(response)