text: import copy
  import os
  from PIL import Image
  
  import torch
  import torch.utils.data
  import torchvision
  
  from pycocotools import mask as coco_mask
  from pycocotools.coco import COCO
  
  import transforms as T
  
  
  def convert_coco_poly_to_mask(segmentations, height, width):
      masks = []
      for polygons in segmentations:
          rles = coco_mask.frPyObjects(polygons, height, width)
          mask = coco_mask.decode(rles)
          if len(mask.shape) < 3:
              mask = mask[..., None]
          mask = torch.as_tensor(mask, dtype=torch.uint8)
          mask = mask.any(dim=2)
          masks.append(mask)
      if masks:
          masks = torch.stack(masks, dim=0)
      else:
          masks = torch.zeros((0, height, width), dtype=torch.uint8)
      return masks
  
  
  class ConvertCocoPolysToMask(object):
      def __init__(self, filter_iscrowd=True):
          self.filter_iscrowd = filter_iscrowd
  
      def __call__(self, image, target):
          w, h = image.size
  
          image_id = target["image_id"]
          image_id = torch.tensor([image_id])
  
          anno = target["annotations"]
  
          if self.filter_iscrowd:
              anno = [obj for obj in anno if obj['iscrowd'] == 0]
  
          boxes = [obj["bbox"] for obj in anno]
          # guard against no boxes via resizing
          boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
          boxes[:, 2:] += boxes[:, :2]
          boxes[:, 0::2].clamp_(min=0, max=w)
          boxes[:, 1::2].clamp_(min=0, max=h)
  
          classes = [obj["category_id"] for obj in anno]
          classes = torch.tensor(classes, dtype=torch.int64)
  
          keypoints = None
          if anno and "keypoints" in anno[0]:
              keypoints = [obj["keypoints"] for obj in anno]
              keypoints = torch.as_tensor(keypoints, dtype=torch.float32)
              num_keypoints = keypoints.shape[0]
              if num_keypoints:
                  keypoints = keypoints.view(num_keypoints, -1, 3)
  
          keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
          boxes = boxes[keep]
          classes = classes[keep]
  
          target = {}
          target["boxes"] = boxes
          target["labels"] = classes
          target["image_id"] = image_id
  
          # for conversion to coco api
          area = torch.tensor([obj["area"] for obj in anno])
          iscrowd = torch.tensor([obj["iscrowd"] for obj in anno])
          target["area"] = area
          target["iscrowd"] = iscrowd
  
          return image, target
  
  
  def _coco_remove_images_without_annotations(dataset, cat_list=None):
      def _has_only_empty_bbox(anno):
          return all(any(o <= 1 for o in obj["bbox"][2:]) for obj in anno)
  
      def _count_visible_keypoints(anno):
          return sum(sum(1 for v in ann["keypoints"][2::3] if v > 0) for ann in anno)
  
      min_keypoints_per_image = 10
  
      def _has_valid_annotation(anno):
          # if it's empty, there is no annotation
          if len(anno) == 0:
              return False
          # if all boxes have close to zero area, there is no annotation
          if _has_only_empty_bbox(anno):
              return False
          # keypoints task have a slight different critera for considering
          # if an annotation is valid
          if "keypoints" not in anno[0]:
              return True
          # for keypoint detection tasks, only consider valid images those
          # containing at least min_keypoints_per_image
          if _count_visible_keypoints(anno) >= min_keypoints_per_image:
              return True
          return False
  
      assert isinstance(dataset, torchvision.datasets.CocoDetection)
      ids = []
      for ds_idx, img_id in enumerate(dataset.ids):
          ann_ids = dataset.coco.getAnnIds(imgIds=img_id, iscrowd=None)
          anno = dataset.coco.loadAnns(ann_ids)
          if cat_list:
              anno = [obj for obj in anno if obj["category_id"] in cat_list]
          if _has_valid_annotation(anno):
              ids.append(ds_idx)
  
      dataset = torch.utils.data.Subset(dataset, ids)
      return dataset
  
  
  def convert_to_coco_api(ds):
      coco_ds = COCO()
      # annotation IDs need to start at 1, not 0, see torchvision issue #1530
      ann_id = 1
      dataset = {'images': [], 'categories': [], 'annotations': []}
      categories = set()
      for img_idx in range(len(ds)):
          # find better way to get target
          # targets = ds.get_annotations(img_idx)
          img, targets = ds[img_idx]
          image_id = targets["image_id"].item()
          img_dict = {}
          img_dict['id'] = image_id
          img_dict['height'] = img.shape[-2]
          img_dict['width'] = img.shape[-1]
          dataset['images'].append(img_dict)
          bboxes = targets["boxes"]
          bboxes[:, 2:] -= bboxes[:, :2]
          bboxes = bboxes.tolist()
          labels = targets['labels'].tolist()
          areas = targets['area'].tolist()
          iscrowd = targets['iscrowd'].tolist()
          num_objs = len(bboxes)
          for i in range(num_objs):
              ann = {}
              ann['image_id'] = image_id
              ann['bbox'] = bboxes[i]
              ann['category_id'] = labels[i]
              categories.add(labels[i])
              ann['area'] = areas[i]
              ann['iscrowd'] = iscrowd[i]
              ann['id'] = ann_id
              dataset['annotations'].append(ann)
              ann_id += 1
      dataset['categories'] = [{'id': i} for i in sorted(categories)]
      coco_ds.dataset = dataset
      coco_ds.createIndex()
      return coco_ds
  
  
  def get_coco_api_from_dataset(dataset):
      for _ in range(10):
          if isinstance(dataset, torchvision.datasets.CocoDetection):
              break
          if isinstance(dataset, torch.utils.data.Subset):
              dataset = dataset.dataset
      if isinstance(dataset, torchvision.datasets.CocoDetection):
          return dataset.coco
      return convert_to_coco_api(dataset)
  
  
  class CocoDetection(torchvision.datasets.CocoDetection):
      def __init__(self, img_folder, ann_file, transforms):
          super(CocoDetection, self).__init__(img_folder, ann_file)
          self._transforms = transforms
  
      def __getitem__(self, idx):
          img, target = super(CocoDetection, self).__getitem__(idx)
          image_id = self.ids[idx]
          target = dict(image_id=image_id, annotations=target)
          if self._transforms is not None:
              img, target = self._transforms(img, target)
          return img, target
  
  
  def get_coco(name, root, image_set, transforms, mode='instances'):
      anno_file_template = "{}_{}2017.json"
      PATHS = {
          "train": ("train2017", os.path.join("annotations", anno_file_template.format(mode, "train"))),
          "val": ("val2017", os.path.join("annotations", anno_file_template.format(mode, "val"))),
      }
  
      t = [ConvertCocoPolysToMask(filter_iscrowd=True)]
  
      if transforms is not None:
          t.append(transforms)
      transforms = T.Compose(t)
  
      img_folder, ann_file = PATHS[image_set]
      img_folder = os.path.join(root, img_folder)
      ann_file = os.path.join(root, ann_file)
  
      dataset = CocoDetection(img_folder, ann_file, transforms=transforms)
  
      if image_set == "train":
          dataset = _coco_remove_images_without_annotations(dataset)
  
      return dataset
  
  
  def get_openimages(name, root, image_set, transforms):
      PATHS = {
          "train": os.path.join(root, "train"),
          "val":   os.path.join(root, "validation"),
      }
  
      t = [ConvertCocoPolysToMask(filter_iscrowd=False)]
  
      if transforms is not None:
          t.append(transforms)
      transforms = T.Compose(t)
  
      img_folder = os.path.join(PATHS[image_set], "data")
      ann_file = os.path.join(PATHS[image_set], "labels", f"{name}.json")
  
      dataset = CocoDetection(img_folder, ann_file, transforms=transforms)
  
      return dataset
  