text: """
  Copyright 2020 The OneFlow Authors. All rights reserved.
  
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  
      http://www.apache.org/licenses/LICENSE-2.0
  
  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
  """
  
  import math
  import os
  import unittest
  
  import cv2
  import numpy as np
  
  import oneflow.compatible.single_client.unittest
  from oneflow.compatible import single_client as flow
  
  VERBOSE = False
  coco_dict = dict()
  
  
  def _coco(anno_file):
      global coco_dict
      if anno_file not in coco_dict:
          from pycocotools.coco import COCO
  
          coco_dict[anno_file] = COCO(anno_file)
      return coco_dict[anno_file]
  
  
  def _make_coco_data_load_fn(
      anno_file,
      image_dir,
      nthread,
      batch_size,
      stride_partition,
      shuffle_after_epoch,
      ret_image_id_only=False,
  ):
      flow.clear_default_session()
      flow.config.cpu_device_num(4)
      func_config = flow.FunctionConfig()
      func_config.default_data_type(flow.float)
      func_config.default_logical_view(flow.scope.consistent_view())
  
      @flow.global_function(function_config=func_config)
      def coco_load_fn():
          with flow.scope.placement("cpu", "0:0-{}".format(nthread - 1)):
              (
                  image,
                  image_id,
                  image_size,
                  gt_bbox,
                  gt_label,
                  gt_segm,
                  gt_segm_index,
              ) = flow.data.coco_reader(
                  annotation_file=anno_file,
                  image_dir=image_dir,
                  batch_size=batch_size,
                  shuffle=shuffle_after_epoch,
                  stride_partition=stride_partition,
                  name="COCOReader",
              )
              if ret_image_id_only:
                  return image_id
              decoded_image = flow.image_decode(image, dtype=flow.float)
              image_list = flow.tensor_buffer_to_tensor_list(
                  decoded_image, shape=(800, 1333, 3), dtype=flow.float
              )
              bbox_list = flow.tensor_buffer_to_tensor_list(
                  gt_bbox, shape=(128, 4), dtype=flow.float
              )
              label_list = flow.tensor_buffer_to_tensor_list(
                  gt_label, shape=(128,), dtype=flow.int32
              )
              segm_list = flow.tensor_buffer_to_tensor_list(
                  gt_segm, shape=(1024, 2), dtype=flow.float
              )
              segm_index_list = flow.tensor_buffer_to_tensor_list(
                  gt_segm_index, shape=(1024, 3), dtype=flow.int32
              )
          return (
              image_id,
              image_size,
              image_list,
              bbox_list,
              label_list,
              segm_list,
              segm_index_list,
          )
  
      return coco_load_fn
  
  
  def _get_coco_image_samples(anno_file, image_dir, image_ids):
      coco = _coco(anno_file)
      category_id_to_contiguous_id_map = _get_category_id_to_contiguous_id_map(coco)
      (image, image_size) = _read_images_with_cv(coco, image_dir, image_ids)
      bbox = _read_bbox(coco, image_ids)
      label = _read_label(coco, image_ids, category_id_to_contiguous_id_map)
      img_segm_poly_list = _read_segm_poly(coco, image_ids)
      (poly, poly_index) = _segm_poly_list_to_tensor(img_segm_poly_list)
      samples = []
      for (im, ims, b, l, p, pi) in zip(image, image_size, bbox, label, poly, poly_index):
          samples.append(
              dict(image=im, image_size=ims, bbox=b, label=l, poly=p, poly_index=pi)
          )
      return samples
  
  
  def _get_category_id_to_contiguous_id_map(coco):
      return {v: i + 1 for (i, v) in enumerate(coco.getCatIds())}
  
  
  def _read_images_with_cv(coco, image_dir, image_ids):
      image_files = [
          os.path.join(image_dir, coco.imgs[img_id]["file_name"]) for img_id in image_ids
      ]
      image_size = [
          (coco.imgs[img_id]["height"], coco.imgs[img_id]["width"])
          for img_id in image_ids
      ]
      return (
          [cv2.imread(image_file).astype(np.single) for image_file in image_files],
          image_size,
      )
  
  
  def _bbox_convert_from_xywh_to_xyxy(bbox, image_h, image_w):
      (x, y, w, h) = bbox
      (x1, y1) = (x, y)
      x2 = x1 + max(w - 1, 0)
      y2 = y1 + max(h - 1, 0)
      x1 = min(max(x1, 0), image_w - 1)
      y1 = min(max(y1, 0), image_h - 1)
      x2 = min(max(x2, 0), image_w - 1)
      y2 = min(max(y2, 0), image_h - 1)
      if x1 >= x2 or y1 >= y2:
          return None
      return [x1, y1, x2, y2]
  
  
  def _read_bbox(coco, image_ids):
      img_bbox_list = []
      for img_id in image_ids:
          anno_ids = coco.getAnnIds(imgIds=[img_id])
          assert len(anno_ids) > 0, "image with id {} has no anno".format(img_id)
          image_h = coco.imgs[img_id]["height"]
          image_w = coco.imgs[img_id]["width"]
          bbox_list = []
          for anno_id in anno_ids:
              anno = coco.anns[anno_id]
              if anno["iscrowd"] != 0:
                  continue
              bbox = anno["bbox"]
              assert isinstance(bbox, list)
              bbox_ = _bbox_convert_from_xywh_to_xyxy(bbox, image_h, image_w)
              if bbox_ is not None:
                  bbox_list.append(bbox_)
          bbox_array = np.array(bbox_list, dtype=np.single)
          img_bbox_list.append(bbox_array)
      return img_bbox_list
  
  
  def _read_label(coco, image_ids, category_id_to_contiguous_id_map):
      img_label_list = []
      for img_id in image_ids:
          anno_ids = coco.getAnnIds(imgIds=[img_id])
          assert len(anno_ids) > 0, "image with id {} has no anno".format(img_id)
          label_list = []
          for anno_id in anno_ids:
              anno = coco.anns[anno_id]
              if anno["iscrowd"] != 0:
                  continue
              cate_id = anno["category_id"]
              isinstance(cate_id, int)
              label_list.append(category_id_to_contiguous_id_map[cate_id])
          label_array = np.array(label_list, dtype=np.int32)
          img_label_list.append(label_array)
      return img_label_list
  
  
  def _read_segm_poly(coco, image_ids):
      img_segm_poly_list = []
      for img_id in image_ids:
          anno_ids = coco.getAnnIds(imgIds=[img_id])
          assert len(anno_ids) > 0, "img {} has no anno".format(img_id)
          segm_poly_list = []
          for anno_id in anno_ids:
              anno = coco.anns[anno_id]
              if anno["iscrowd"] != 0:
                  continue
              segm = anno["segmentation"]
              assert isinstance(segm, list)
              assert len(segm) > 0, str(len(segm))
              assert all([len(poly) > 0 for poly in segm]), str(
                  [len(poly) for poly in segm]
              )
              segm_poly_list.append(segm)
          img_segm_poly_list.append(segm_poly_list)
      return img_segm_poly_list
  
  
  def _segm_poly_list_to_tensor(img_segm_poly_list):
      poly_array_list = []
      poly_index_array_list = []
      for (img_idx, segm_poly_list) in enumerate(img_segm_poly_list):
          img_poly_elem_list = []
          img_poly_index_list = []
          for (obj_idx, poly_list) in enumerate(segm_poly_list):
              for (poly_idx, poly) in enumerate(poly_list):
                  img_poly_elem_list.extend(poly)
                  for (pt_idx, pt) in enumerate(poly):
                      if pt_idx % 2 == 0:
                          img_poly_index_list.append([pt_idx / 2, poly_idx, obj_idx])
          img_poly_array = np.array(img_poly_elem_list, dtype=np.single).reshape(-1, 2)
          assert img_poly_array.size > 0, segm_poly_list
          poly_array_list.append(img_poly_array)
          img_poly_index_array = np.array(img_poly_index_list, dtype=np.int32)
          assert img_poly_index_array.size > 0, segm_poly_list
          poly_index_array_list.append(img_poly_index_array)
      return (poly_array_list, poly_index_array_list)
  
  
  def _get_coco_sorted_imgs(anno_file):
      coco = _coco(anno_file)
      img_ids = coco.getImgIds()
      img_ids.sort()
      img_info_list = []
      for (i, img_id) in enumerate(img_ids):
          img_h = coco.imgs[img_id]["height"]
          img_w = coco.imgs[img_id]["width"]
          group_id = int(img_h / img_w)
          anno_ids = coco.getAnnIds(imgIds=img_id, iscrowd=None)
          anno = coco.loadAnns(anno_ids)
          if not _has_valid_annotation(anno):
              continue
          img_info_list.append(
              dict(index=i, image_id=img_id, group_id=group_id, anno_len=len(anno_ids))
          )
      return img_info_list
  
  
  def _count_visible_keypoints(anno):
      return sum((sum((1 for v in ann["keypoints"][2::3] if v > 0)) for ann in anno))
  
  
  def _has_only_empty_bbox(anno):
      return all((any((o <= 1 for o in obj["bbox"][2:])) for obj in anno))
  
  
  def _has_valid_annotation(anno):
      if len(anno) == 0:
          return False
      if _has_only_empty_bbox(anno):
          return False
      if "keypoints" not in anno[0]:
          return True
      if _count_visible_keypoints(anno) >= 10:
          return True
      return False
  
  
  class GroupedDistributedSampler(object):
      def __init__(self, shards, batch_size, images, stride_sample, max_iter=3):
          assert batch_size % shards == 0
          self._images = images
          self._shards = shards
          self._shard_size = math.ceil(len(images) / shards)
          self._batch_size = batch_size
          self._batch_size_per_shard = batch_size // shards
          self._stride_sample = stride_sample
          self._max_iter = max_iter
          self._init_sample_idx()
          self._init_group_buckets()
  
      def _init_sample_idx(self):
          if self._stride_sample:
              self._sample_idx = list(range(self._shards))
          else:
              self._sample_idx = [rank * self._shard_size for rank in range(self._shards)]
              self._sample_idx_in_shard = [0 for _ in range(self._shards)]
  
      def _init_group_buckets(self):
          self._group_buckets = [[[] for _ in range(2)] for _ in range(self._shards)]
  
      def __iter__(self):
          for i in range(self._max_iter):
              sample_ids = []
              for rank in range(self._shards):
                  sample_cnt_cur_rank = 0
                  sample_ids_cur_rank = []
                  group_buckets_cur_rank = self._group_buckets[rank]
                  if (
                      len(group_buckets_cur_rank[0]) > 0
                      and len(group_buckets_cur_rank[1]) > 0
                  ):
                      if (
                          group_buckets_cur_rank[0][0]["index"]
                          < group_buckets_cur_rank[1][0]["index"]
                      ):
                          sample = group_buckets_cur_rank[0].pop(0)
                      else:
                          sample = group_buckets_cur_rank[1].pop(0)
                  elif len(group_buckets_cur_rank[0]) > 0:
                      sample = group_buckets_cur_rank[0].pop(0)
                  elif len(group_buckets_cur_rank[1]) > 0:
                      sample = group_buckets_cur_rank[1].pop(0)
                  else:
                      sample = self._next_sample(rank)
                  group_id = sample["group_id"]
                  sample_ids_cur_rank.append(sample["image_id"])
                  sample_cnt_cur_rank += 1
                  while sample_cnt_cur_rank < self._batch_size_per_shard:
                      if len(group_buckets_cur_rank[group_id]) > 0:
                          sample = group_buckets_cur_rank[group_id].pop(0)
                          sample_ids_cur_rank.append(sample["image_id"])
                          sample_cnt_cur_rank += 1
                          continue
                      sample = self._next_sample(rank)
                      if sample["group_id"] == group_id:
                          sample_ids_cur_rank.append(sample["image_id"])
                          sample_cnt_cur_rank += 1
                      else:
                          group_buckets_cur_rank[sample["group_id"]].append(sample)
                  sample_ids.extend(sample_ids_cur_rank)
              yield sample_ids
  
      def _next_sample(self, rank):
          sample = self._images[self._sample_idx[rank]]
          if self._stride_sample:
              self._sample_idx[rank] += self._shards
          else:
              self._sample_idx_in_shard[rank] += 1
              self._sample_idx[rank] += 1
              if self._sample_idx_in_shard[rank] == self._shard_size:
                  self._sample_idx[rank] += (self._shards - 1) * self._shard_size
                  self._sample_idx_in_shard[rank] = 0
          if self._sample_idx[rank] >= len(self._images):
              self._sample_idx[rank] %= len(self._images)
          return sample
  
  
  @unittest.skipIf(True, "skip for now because of single-client tensor_list removed")
  class TestCocoReader(flow.unittest.TestCase):
      def test_coco_reader(test_case, verbose=VERBOSE):
          anno_file = "/dataset/mscoco_2017/annotations/instances_val2017.json"
          image_dir = "/dataset/mscoco_2017/val2017"
          of_coco_load_fn = _make_coco_data_load_fn(
              anno_file, image_dir, 1, 2, True, True
          )
          (
              image_id,
              image_size,
              image,
              bbox,
              label,
              poly,
              poly_index,
          ) = of_coco_load_fn().get()
          image_id = image_id.numpy()
          image_size = image_size.numpy()
          image = image.numpy_lists()
          bbox = bbox.numpy_lists()
          label = label.numpy_lists()
          poly = poly.numpy_lists()
          poly_index = poly_index.numpy_lists()
          samples = _get_coco_image_samples(anno_file, image_dir, image_id)
          for (i, sample) in enumerate(samples):
              if verbose:
                  print(
                      "#{} of label:\n".format(i),
                      label[0][i].squeeze(),
                      type(label[0][i].squeeze()),
                      label[0][i].squeeze().shape,
                  )
                  print(
                      "#{} coco label:\n".format(i),
                      sample["label"],
                      type(sample["label"]),
                      sample["label"].shape,
                  )
              test_case.assertTrue(np.array_equal(image[0][i].squeeze(), sample["image"]))
              test_case.assertTrue(np.array_equal(image_size[i], sample["image_size"]))
              test_case.assertTrue(np.allclose(bbox[0][i].squeeze(), sample["bbox"]))
              cur_label = label[0][i].squeeze()
              if len(cur_label.shape) == 0:
                  cur_label = np.array([cur_label])
              test_case.assertTrue(np.array_equal(cur_label, sample["label"]))
              test_case.assertTrue(np.allclose(poly[0][i].squeeze(), sample["poly"]))
              test_case.assertTrue(
                  np.array_equal(poly_index[0][i].squeeze(), sample["poly_index"])
              )
  
      def test_coco_reader_distributed_stride(test_case, verbose=VERBOSE):
          anno_file = "/dataset/mscoco_2017/annotations/instances_val2017.json"
          image_dir = "/dataset/mscoco_2017/val2017"
          image_info_list = _get_coco_sorted_imgs(anno_file)
          if verbose:
              print("Info of the first 20 images:")
              for (i, image_info) in enumerate(image_info_list[:20]):
                  print(
                      "index: {}, image_id: {}, group_id: {}, anno len: {}".format(
                          i,
                          image_info["image_id"],
                          image_info["group_id"],
                          image_info["anno_len"],
                      )
                  )
          sampler = GroupedDistributedSampler(4, 8, image_info_list, True)
          of_coco_load_fn = _make_coco_data_load_fn(
              anno_file, image_dir, 4, 8, True, False, True
          )
          for (i, sample_ids) in enumerate(sampler):
              image_id = of_coco_load_fn().get().numpy()
              if verbose:
                  print("#{} image_id:".format(i), image_id)
                  print("#{} sample_ids:".format(i), sample_ids)
              test_case.assertTrue(np.array_equal(image_id, sample_ids))
  
      def test_coco_reader_distributed_contiguous(test_case, verbose=VERBOSE):
          anno_file = "/dataset/mscoco_2017/annotations/instances_val2017.json"
          image_dir = "/dataset/mscoco_2017/val2017"
          image_info_list = _get_coco_sorted_imgs(anno_file)
          sampler = GroupedDistributedSampler(4, 8, image_info_list, False)
          of_coco_load_fn = _make_coco_data_load_fn(
              anno_file, image_dir, 4, 8, False, False, True
          )
          for (i, sample_ids) in enumerate(sampler):
              image_id = of_coco_load_fn().get().numpy()
              if verbose:
                  print("#{} image_id:".format(i), image_id)
                  print("#{} sample_ids:".format(i), sample_ids)
              test_case.assertTrue(np.array_equal(image_id, sample_ids))
  
  
  if __name__ == "__main__":
      unittest.main()
  