text: #!/usr/bin/env python3
  # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
  
  
  import contextlib
  import itertools
  import json
  import os
  import uuid
  
  from d2go.data.datasets import register_dataset_split
  from detectron2.data import DatasetCatalog, MetadataCatalog
  from mobile_cv.common.misc.file_utils import make_temp_directory
  from PIL import Image
  
  
  IM_DIR = "image_directory"
  ANN_FN = "annotation_file"
  
  
  def create_toy_dataset(
      image_generator, num_images, num_classes=-1, num_keypoints=0, is_rotated=False
  ):
      """ given image_generator, create a dataset with toy annotations and catagories """
      categories = []
      images = []
      annotations = []
      meta_data = {}
  
      if num_classes == -1:
          num_classes = num_images
  
      for i in range(num_images):
          image_generator.prepare_image(i)
          image_dict = image_generator.get_image_dict(i)
          width = image_dict["width"]
          height = image_dict["height"]
          images.append(image_dict)
  
          if i < num_classes:
              categories.append({"name": "class_{}".format(i), "id": i})
  
          bbox = (
              [width / 4, height / 4, width / 2, height / 2]  # XYWH_ABS
              if not is_rotated
              else [width / 2, height / 2, width / 2, height / 2, 45] # cXcYWHO_ABS
          )
  
          annotations.append(
              {
                  "image_id": i,
                  "category_id": i % num_classes,
                  "id": i,
                  "bbox": bbox,
                  "keypoints": list(
                      itertools.chain.from_iterable(
                          [
                              (
                                  float(idx) / width / 2 + width / 4,
                                  float(idx) / height / 2 + height / 4,
                                  1,
                              )
                              for idx in range(num_keypoints)
                          ]
                      )
                  ),
                  "area": width * height,
                  "iscrowd": 0,
                  "ignore": 0,
                  "segmentation": [],
              }
          )
  
      if num_keypoints > 0:
          keypoint_names = [f"kp_{idx}" for idx in range(num_keypoints)]
          meta_data.update({"keypoint_names": keypoint_names, "keypoint_flip_map": ()})
  
      return (
          {"categories": categories, "images": images, "annotations": annotations},
          meta_data,
      )
  
  
  @contextlib.contextmanager
  def register_toy_dataset(
      dataset_name, image_generator, num_images, num_classes=-1, num_keypoints=0
  ):
      json_dataset, meta_data = create_toy_dataset(
          image_generator,
          num_images=num_images,
          num_classes=num_classes,
          num_keypoints=num_keypoints,
      )
  
      with make_temp_directory("detectron2go_tmp_dataset") as tmp_dir:
          json_file = os.path.join(tmp_dir, "{}.json".format(dataset_name))
          with open(json_file, "w") as f:
              json.dump(json_dataset, f)
  
          split_dict = {
              IM_DIR: image_generator.get_image_dir(),
              ANN_FN: json_file,
              "meta_data": meta_data,
          }
          register_dataset_split(dataset_name, split_dict)
  
          try:
              yield
          finally:
              DatasetCatalog.remove(dataset_name)
              MetadataCatalog.remove(dataset_name)
  
  
  def create_local_dataset(
      out_dir,
      num_images,
      image_width,
      image_height,
      num_classes=-1,
      num_keypoints=0,
      is_rotated=False,
  ):
      dataset_name = "_test_ds_" + str(uuid.uuid4())
  
      img_gen = LocalImageGenerator(out_dir, image_width, image_height)
      json_dataset, meta_data = create_toy_dataset(
          img_gen,
          num_images=num_images,
          num_classes=num_classes,
          num_keypoints=num_keypoints,
      )
      json_file = os.path.join(out_dir, "{}.json".format(dataset_name))
      with open(json_file, "w") as f:
          json.dump(json_dataset, f)
  
      split_dict = {
          IM_DIR: img_gen.get_image_dir(),
          ANN_FN: json_file,
          "meta_data": meta_data,
      }
      if is_rotated:
          split_dict['evaluator_type'] = "rotated_coco"
      register_dataset_split(dataset_name, split_dict)
  
      return dataset_name
  
  
  class LocalImageGenerator:
      def __init__(self, image_dir, width, height):
          self._width = width
          self._height = height
          self._image_dir = image_dir
  
      def get_image_dir(self):
          return self._image_dir
  
      def get_image_dict(self, i):
          return {
              "file_name": "{}.jpg".format(i),
              "width": self._width,
              "height": self._height,
              "id": i,
          }
  
      def prepare_image(self, i):
          image = Image.new("RGB", (self._width, self._height))
          image.save(os.path.join(self._image_dir, self.get_image_dict(i)["file_name"]))
  