code: # Ultralytics YOLO 🚀, AGPL-3.0 license
  
  import getpass
  from typing import List
  
  import cv2
  import numpy as np
  import pandas as pd
  
  from ultralytics.data.augment import LetterBox
  from ultralytics.utils import LOGGER as logger
  from ultralytics.utils import SETTINGS
  from ultralytics.utils.checks import check_requirements
  from ultralytics.utils.ops import xyxy2xywh
  from ultralytics.utils.plotting import plot_images
  
  
  def get_table_schema(vector_size):
      """Extracts and returns the schema of a database table."""
      from lancedb.pydantic import LanceModel, Vector
  
      class Schema(LanceModel):
          im_file: str
          labels: List[str]
          cls: List[int]
          bboxes: List[List[float]]
          masks: List[List[List[int]]]
          keypoints: List[List[List[float]]]
          vector: Vector(vector_size)
  
      return Schema
  
  
  def get_sim_index_schema():
      """Returns a LanceModel schema for a database table with specified vector size."""
      from lancedb.pydantic import LanceModel
  
      class Schema(LanceModel):
          idx: int
          im_file: str
          count: int
          sim_im_files: List[str]
  
      return Schema
  
  
  def sanitize_batch(batch, dataset_info):
      """Sanitizes input batch for inference, ensuring correct format and dimensions."""
      batch["cls"] = batch["cls"].flatten().int().tolist()
      box_cls_pair = sorted(zip(batch["bboxes"].tolist(), batch["cls"]), key=lambda x: x[1])
      batch["bboxes"] = [box for box, _ in box_cls_pair]
      batch["cls"] = [cls for _, cls in box_cls_pair]
      batch["labels"] = [dataset_info["names"][i] for i in batch["cls"]]
      batch["masks"] = batch["masks"].tolist() if "masks" in batch else [[[]]]
      batch["keypoints"] = batch["keypoints"].tolist() if "keypoints" in batch else [[[]]]
      return batch
  
  
  def plot_query_result(similar_set, plot_labels=True):
      """
      Plot images from the similar set.
  
      Args:
          similar_set (list): Pyarrow or pandas object containing the similar data points
          plot_labels (bool): Whether to plot labels or not
      """
      similar_set = (
          similar_set.to_dict(orient="list") if isinstance(similar_set, pd.DataFrame) else similar_set.to_pydict()
      )
      empty_masks = [[[]]]
      empty_boxes = [[]]
      images = similar_set.get("im_file", [])
      bboxes = similar_set.get("bboxes", []) if similar_set.get("bboxes") is not empty_boxes else []
      masks = similar_set.get("masks") if similar_set.get("masks")[0] != empty_masks else []
      kpts = similar_set.get("keypoints") if similar_set.get("keypoints")[0] != empty_masks else []
      cls = similar_set.get("cls", [])
  
      plot_size = 640
      imgs, batch_idx, plot_boxes, plot_masks, plot_kpts = [], [], [], [], []
      for i, imf in enumerate(images):
          im = cv2.imread(imf)
          im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
          h, w = im.shape[:2]
          r = min(plot_size / h, plot_size / w)
          imgs.append(LetterBox(plot_size, center=False)(image=im).transpose(2, 0, 1))
          if plot_labels:
              if len(bboxes) > i and len(bboxes[i]) > 0:
                  box = np.array(bboxes[i], dtype=np.float32)
                  box[:, [0, 2]] *= r
                  box[:, [1, 3]] *= r
                  plot_boxes.append(box)
              if len(masks) > i and len(masks[i]) > 0:
                  mask = np.array(masks[i], dtype=np.uint8)[0]
                  plot_masks.append(LetterBox(plot_size, center=False)(image=mask))
              if len(kpts) > i and kpts[i] is not None:
                  kpt = np.array(kpts[i], dtype=np.float32)
                  kpt[:, :, :2] *= r
                  plot_kpts.append(kpt)
          batch_idx.append(np.ones(len(np.array(bboxes[i], dtype=np.float32))) * i)
      imgs = np.stack(imgs, axis=0)
      masks = np.stack(plot_masks, axis=0) if len(plot_masks) > 0 else np.zeros(0, dtype=np.uint8)
      kpts = np.concatenate(plot_kpts, axis=0) if len(plot_kpts) > 0 else np.zeros((0, 51), dtype=np.float32)
      boxes = xyxy2xywh(np.concatenate(plot_boxes, axis=0)) if len(plot_boxes) > 0 else np.zeros(0, dtype=np.float32)
      batch_idx = np.concatenate(batch_idx, axis=0)
      cls = np.concatenate([np.array(c, dtype=np.int32) for c in cls], axis=0)
  
      return plot_images(
          imgs, batch_idx, cls, bboxes=boxes, masks=masks, kpts=kpts, max_subplots=len(images), save=False, threaded=False
      )
  
  
  def prompt_sql_query(query):
      """Plots images with optional labels from a similar data set."""
      check_requirements("openai>=1.6.1")
      from openai import OpenAI
  
      if not SETTINGS["openai_api_key"]:
          logger.warning("OpenAI API key not found in settings. Please enter your API key below.")
          openai_api_key = getpass.getpass("OpenAI API key: ")
          SETTINGS.update({"openai_api_key": openai_api_key})
      openai = OpenAI(api_key=SETTINGS["openai_api_key"])
  
      messages = [
          {
              "role": "system",
              "content": """
                  You are a helpful data scientist proficient in SQL. You need to output exactly one SQL query based on
                  the following schema and a user request. You only need to output the format with fixed selection
                  statement that selects everything from "'table'", like `SELECT * from 'table'`
  
                  Schema:
                  im_file: string not null
                  labels: list<item: string> not null
                  child 0, item: string
                  cls: list<item: int64> not null
                  child 0, item: int64
                  bboxes: list<item: list<item: double>> not null
                  child 0, item: list<item: double>
                      child 0, item: double
                  masks: list<item: list<item: list<item: int64>>> not null
                  child 0, item: list<item: list<item: int64>>
                      child 0, item: list<item: int64>
                          child 0, item: int64
                  keypoints: list<item: list<item: list<item: double>>> not null
                  child 0, item: list<item: list<item: double>>
                      child 0, item: list<item: double>
                          child 0, item: double
                  vector: fixed_size_list<item: float>[256] not null
                  child 0, item: float
  
                  Some details about the schema:
                  - the "labels" column contains the string values like 'person' and 'dog' for the respective objects
                      in each image
                  - the "cls" column contains the integer values on these classes that map them the labels
  
                  Example of a correct query:
                  request - Get all data points that contain 2 or more people and at least one dog
                  correct query-
                  SELECT * FROM 'table' WHERE  ARRAY_LENGTH(cls) >= 2  AND ARRAY_LENGTH(FILTER(labels, x -> x = 'person')) >= 2  AND ARRAY_LENGTH(FILTER(labels, x -> x = 'dog')) >= 1;
               """,
          },
          {"role": "user", "content": f"{query}"},
      ]
  
      response = openai.chat.completions.create(model="gpt-3.5-turbo", messages=messages)
      return response.choices[0].message.content
  
apis: lancedb.pydantic.Vector
extract_api: [((3694, 3716), 'numpy.stack', 'np.stack', (['imgs'], {'axis': '(0)'}), '(imgs, axis=0)\n', (3702, 3716), True, 'import numpy as np\n'), ((4054, 4087), 'numpy.concatenate', 'np.concatenate', (['batch_idx'], {'axis': '(0)'}), '(batch_idx, axis=0)\n', (4068, 4087), True, 'import numpy as np\n'), ((4421, 4456), 'ultralytics.utils.checks.check_requirements', 'check_requirements', (['"""openai>=1.6.1"""'], {}), "('openai>=1.6.1')\n", (4439, 4456), False, 'from ultralytics.utils.checks import check_requirements\n'), ((4758, 4800), 'openai.OpenAI', 'OpenAI', ([], {'api_key': "SETTINGS['openai_api_key']"}), "(api_key=SETTINGS['openai_api_key'])\n", (4764, 4800), False, 'from openai import OpenAI\n'), ((798, 817), 'lancedb.pydantic.Vector', 'Vector', (['vector_size'], {}), '(vector_size)\n', (804, 817), False, 'from lancedb.pydantic import LanceModel, Vector\n'), ((2746, 2761), 'cv2.imread', 'cv2.imread', (['imf'], {}), '(imf)\n', (2756, 2761), False, 'import cv2\n'), ((2775, 2810), 'cv2.cvtColor', 'cv2.cvtColor', (['im', 'cv2.COLOR_BGR2RGB'], {}), '(im, cv2.COLOR_BGR2RGB)\n', (2787, 2810), False, 'import cv2\n'), ((3729, 3757), 'numpy.stack', 'np.stack', (['plot_masks'], {'axis': '(0)'}), '(plot_masks, axis=0)\n', (3737, 3757), True, 'import numpy as np\n'), ((3786, 3813), 'numpy.zeros', 'np.zeros', (['(0)'], {'dtype': 'np.uint8'}), '(0, dtype=np.uint8)\n', (3794, 3813), True, 'import numpy as np\n'), ((3825, 3858), 'numpy.concatenate', 'np.concatenate', (['plot_kpts'], {'axis': '(0)'}), '(plot_kpts, axis=0)\n', (3839, 3858), True, 'import numpy as np\n'), ((3886, 3921), 'numpy.zeros', 'np.zeros', (['(0, 51)'], {'dtype': 'np.float32'}), '((0, 51), dtype=np.float32)\n', (3894, 3921), True, 'import numpy as np\n'), ((4008, 4037), 'numpy.zeros', 'np.zeros', (['(0)'], {'dtype': 'np.float32'}), '(0, dtype=np.float32)\n', (4016, 4037), True, 'import numpy as np\n'), ((4535, 4628), 'ultralytics.utils.LOGGER.warning', 'logger.warning', (['"""OpenAI API key not found in settings. Please enter your API key below."""'], {}), "(\n    'OpenAI API key not found in settings. Please enter your API key below.')\n", (4549, 4628), True, 'from ultralytics.utils import LOGGER as logger\n'), ((4649, 4684), 'getpass.getpass', 'getpass.getpass', (['"""OpenAI API key: """'], {}), "('OpenAI API key: ')\n", (4664, 4684), False, 'import getpass\n'), ((4693, 4744), 'ultralytics.utils.SETTINGS.update', 'SETTINGS.update', (["{'openai_api_key': openai_api_key}"], {}), "({'openai_api_key': openai_api_key})\n", (4708, 4744), False, 'from ultralytics.utils import SETTINGS\n'), ((3944, 3978), 'numpy.concatenate', 'np.concatenate', (['plot_boxes'], {'axis': '(0)'}), '(plot_boxes, axis=0)\n', (3958, 3978), True, 'import numpy as np\n'), ((4114, 4141), 'numpy.array', 'np.array', (['c'], {'dtype': 'np.int32'}), '(c, dtype=np.int32)\n', (4122, 4141), True, 'import numpy as np\n'), ((3071, 3108), 'numpy.array', 'np.array', (['bboxes[i]'], {'dtype': 'np.float32'}), '(bboxes[i], dtype=np.float32)\n', (3079, 3108), True, 'import numpy as np\n'), ((3492, 3527), 'numpy.array', 'np.array', (['kpts[i]'], {'dtype': 'np.float32'}), '(kpts[i], dtype=np.float32)\n', (3500, 3527), True, 'import numpy as np\n'), ((3296, 3330), 'numpy.array', 'np.array', (['masks[i]'], {'dtype': 'np.uint8'}), '(masks[i], dtype=np.uint8)\n', (3304, 3330), True, 'import numpy as np\n'), ((2905, 2939), 'ultralytics.data.augment.LetterBox', 'LetterBox', (['plot_size'], {'center': '(False)'}), '(plot_size, center=False)\n', (2914, 2939), False, 'from ultralytics.data.augment import LetterBox\n'), ((3368, 3402), 'ultralytics.data.augment.LetterBox', 'LetterBox', (['plot_size'], {'center': '(False)'}), '(plot_size, center=False)\n', (3377, 3402), False, 'from ultralytics.data.augment import LetterBox\n'), ((3638, 3675), 'numpy.array', 'np.array', (['bboxes[i]'], {'dtype': 'np.float32'}), '(bboxes[i], dtype=np.float32)\n', (3646, 3675), True, 'import numpy as np\n')]