LMDX: Cloud Vision API x LLMで画像から構造化データ作成

import io
from google.cloud import vision

client = vision.ImageAnnotatorClient()

# TODO: fileを定義
with io.open(file, 'rb') as image_file:
      content = image_file.read()
  image = vision.Image(content=content)
  res = client.document_text_detection(
      image=image,
      image_context={'language_hints': ['ja']})
  ocr_result = res

Vertex = namedtuple('Vertex', ['x', 'y'])
Sentence = namedtuple('Sentence', ['text', 'bounding_box', 'quantized_center'])

def generate_document(full_text_annotation):
  def main(full_text_annotation):
    sentences = extract_sentences_from_full_text_annotation(full_text_annotation)
    # ここが大事
    document = "".join(f"{sentence.text} {sentence.quantized_center.x}|{sentence.quantized_center.y}\n" for sentence in sentences)
    return document

  def quantize(value, min_value, max_value, bins):
      bin_size = (max_value - min_value) / bins
      return int((value - min_value) // bin_size)

  def extract_sentences_from_full_text_annotation(full_text_annotation, bins = 100):
      sentences = []
      
      for page in full_text_annotation.pages:
          max_x = page.width
          max_y = page.height
          for block in page.blocks:
              for paragraph in block.paragraphs:
                  paragraph_text = ""
                  for word in paragraph.words:
                      word_text = "".join(symbol.text for symbol in word.symbols)
                      paragraph_text += word_text + " "

                  min_x = 0
                  min_y = 0
                  vertices = paragraph.bounding_box.vertices
                  center_x = sum(vertex.x for vertex in vertices) / 4
                  center_y = sum(vertex.y for vertex in vertices) / 4
                  
                  quantized_center_x = quantize(center_x, min_x, max_x, bins)
                  quantized_center_y = quantize(center_y, min_y, max_y, bins)
                  quantized_center = Vertex(quantized_center_x, quantized_center_y)

                  sentences.append(Sentence(paragraph_text, paragraph.bounding_box, quantized_center))
      return sentences

  return main(full_text_annotation)

def generate_task():
  task = "From the document, extract the text values and tags of the following entities:\n"
  # 適宜、フォーマット変える。
  task += '{ "name": "", "address": "" }\n'
  return task

def generate_prompt(full_text_annotation):
  document = generate_document(full_text_annotation)
  task = generate_task()
  prompt = f"<Document>\n{document}</Document>\n"
  prompt += f"<Task>\n{task}</Task>\n"
  prompt += "<Extraction>\n"
  return prompt

import google.generativeai as genai

model =  "models/gemini-pro"
client = genai.GenerativeModel(model)

def generate_text(prompt):
  response = client.generate_content(prompt)
  return response.text