text: # Copyright (c) Facebook, Inc. and its affiliates.
  
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  import time
  from typing import Callable, List
  
  import torch
  import torchvision.transforms as transforms
  from PIL import Image
  from vissl.data.ssl_transforms.img_cv_color_distortion import (
      ColorDistortionSettings,
      ImgOpenCVColorDistortion,
  )
  from vissl.data.ssl_transforms.img_pil_color_distortion import ImgPilColorDistortion
  from vissl.data.ssl_transforms.img_pil_gaussian_blur import ImgPilGaussianBlur
  from vissl.data.ssl_transforms.img_pil_random_photometric import ImgPilRandomPhotometric
  from vissl.data.ssl_transforms.img_pil_to_lab_tensor import ImgPil2LabTensor
  from vissl.data.ssl_transforms.img_pil_to_multicrop import ImgPilToMultiCrop
  from vissl.data.ssl_transforms.img_pil_to_patches_and_image import (
      ImgPilToPatchesAndImage,
  )
  from vissl.data.ssl_transforms.img_pil_to_tensor import ImgToTensor
  from vissl.data.ssl_transforms.img_rotate_pil import ImgRotatePil
  
  
  """
  Small util to test the throughput of the existing transforms on a cpu host
  """
  
  RAND_TENSOR = (torch.rand((224, 224, 3)) * 255).to(dtype=torch.uint8)
  RAND_PIL = Image.fromarray(RAND_TENSOR.numpy())
  ITERATIONS = 1000
  N_QUEUES = torch.multiprocessing.cpu_count()  # Simulate the load of N dataloader queues
  
  
  # ---- Tooling -----
  def benchmark(transform: Callable, title: str, requires_pil: bool = False) -> float:
      """Given a transform, simulate a real-world load by creating multiple workers"""
  
      #  Test the multi-queue speed
      global benchmark_results
      benchmark_results = []
      pool = torch.multiprocessing.Pool(N_QUEUES)
      pool.starmap_async(
          load_one_queue, [(transform, requires_pil)] * N_QUEUES, callback=store_result
      ).get()
  
      # Test the mono-queue speed
      fps_mono = load_one_queue(transform, requires_pil)
  
      print(
          "{: <25}: {: >10.1f} fps single queue, {: >10.1f} fps multi queue".format(
              title, fps_mono, sum(benchmark_results) / N_QUEUES
          )
      )
  
  
  def store_result(result: List[float]) -> None:
      global benchmark_results
      benchmark_results.extend(result)
  
  
  def load_one_queue(transform: Callable, requires_pil: bool = False) -> float:
      """Run a given transform repeatedly to simulate a single dataloader worker"""
      start = time.time()
      for _ in range(ITERATIONS):
          transform(RAND_PIL) if requires_pil else transform(RAND_TENSOR)
  
      return ITERATIONS / (time.time() - start)
  
  
  # ---- Transforms to be benchmarked -----
  def testBlur():
      transform = ImgPilGaussianBlur(p=0.5, radius_min=0.1, radius_max=2.0)
      benchmark(transform, "Blur", requires_pil=True)
  
  
  def testImPil2Lab():
      benchmark(ImgPil2LabTensor, "PIL2Lab")
  
  
  def testColorDistort():
      transform = ImgPilColorDistortion(strength=0.5)
      benchmark(transform, "Color distort - PIL", requires_pil=True)
  
  
  def testOpenCVColorDistort():
      transform = ImgOpenCVColorDistortion(ColorDistortionSettings.defaults())
      benchmark(transform, "Color distort - OpenCV", requires_pil=True)
  
  
  def testToTensor():
      transform = transforms.ToTensor()
      benchmark(transform, "ToTensor", requires_pil=True)
  
  
  def testImgToTensor():
      benchmark(ImgToTensor(), "ImgToTensor", requires_pil=True)
  
  
  def testImgPatch():
      transform = ImgPilToPatchesAndImage()
      benchmark(transform, "Patches", requires_pil=True)
  
  
  def testImgRotate():
      transform = ImgRotatePil(num_angles=4, num_rotations_per_img=1)
      benchmark(transform, "ImgRotate", requires_pil=True)
  
  
  def testImgPilToMulticrop():
      transform = ImgPilToMultiCrop(
          total_num_crops=8,
          num_crops=[2, 6],
          size_crops=[224, 96],
          crop_scales=[[0.14, 1], [0.05, 0.14]],
      )
  
      benchmark(transform, "MultiCrop", requires_pil=True)
  
  
  def testImgPilRandomPhotometric():
      transform = ImgPilRandomPhotometric(p=0.5)  # reflect a real world use case
      benchmark(transform, "Photometric", requires_pil=True)
  
  
  if __name__ == "__main__":
      # Run the transforms and print out an average processing speed
      print("\n")
      print(f"*** Using {N_QUEUES} {'queues' if N_QUEUES >1 else 'queue'} ***")
      testBlur()
      testColorDistort()
      testOpenCVColorDistort()
      testToTensor()
      testImgToTensor()
      testImPil2Lab()
      testImgPatch()
      testImgRotate()
      testImgPilToMulticrop()
      testImgPilRandomPhotometric()
      print("\n")
  