# coding=utf-8
# Copyright 2020 The TensorFlow Datasets Authors.
#
# 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.

# Lint as: python3
"""MS Coco Dataset."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import json
import os

from absl import logging
import tensorflow.compat.v2 as tf

import tensorflow_datasets.public_api as tfds

_CITATION = """\
@article{DBLP:journals/corr/LinMBHPRDZ14,
  author    = {Tsung{-}Yi Lin and
               Michael Maire and
               Serge J. Belongie and
               Lubomir D. Bourdev and
               Ross B. Girshick and
               James Hays and
               Pietro Perona and
               Deva Ramanan and
               Piotr Doll{\'{a}}r and
               C. Lawrence Zitnick},
  title     = {Microsoft {COCO:} Common Objects in Context},
  journal   = {CoRR},
  volume    = {abs/1405.0312},
  year      = {2014},
  url       = {http://arxiv.org/abs/1405.0312},
  archivePrefix = {arXiv},
  eprint    = {1405.0312},
  timestamp = {Mon, 13 Aug 2018 16:48:13 +0200},
  biburl    = {https://dblp.org/rec/bib/journals/corr/LinMBHPRDZ14},
  bibsource = {dblp computer science bibliography, https://dblp.org}
}
"""

DESCRIPTION = """COCO is a large-scale object detection, segmentation, and
captioning dataset. This version contains images, bounding boxes "
and labels for the {year} version.
Note:
 * Some images from the train and validation sets don't have annotations.
 * Coco 2014 and 2017 uses the same images, but different train/val/test splits
 * The test split don't have any annotations (only images).
 * Coco defines 91 classes but the data only uses 80 classes.
 * Panotptic annotations defines defines 200 classes but only uses 133.
"""


Split = collections.namedtuple(
    'Split', ['name', 'images', 'annotations', 'annotation_type'])


class AnnotationType(object):
  """Enum of the annotation format types.

  Splits are annotated with different formats.
  """
  BBOXES = 'bboxes'
  PANOPTIC = 'panoptic'
  NONE = 'none'


class CocoConfig(tfds.core.BuilderConfig):
  """BuilderConfig for CocoConfig."""

  def __init__(
      self,
      splits=None,
      has_panoptic=False,
      **kwargs):
    super(CocoConfig, self).__init__(
        version=tfds.core.Version('1.1.0'),
        **kwargs)
    self.splits = splits
    self.has_panoptic = has_panoptic


class Coco(tfds.core.GeneratorBasedBuilder):
  """Base MS Coco dataset."""

  BUILDER_CONFIGS = [
      CocoConfig(
          name='2014',
          description=DESCRIPTION.format(year=2014),
          splits=[
              Split(
                  name=tfds.Split.TRAIN,
                  images='train2014',
                  annotations='annotations_trainval2014',
                  annotation_type=AnnotationType.BBOXES,
              ),
              Split(
                  name=tfds.Split.VALIDATION,
                  images='val2014',
                  annotations='annotations_trainval2014',
                  annotation_type=AnnotationType.BBOXES,
              ),
              Split(
                  name=tfds.Split.TEST,
                  images='test2014',
                  annotations='image_info_test2014',
                  annotation_type=AnnotationType.NONE,
              ),
              # Coco2014 contains an extra test split
              Split(
                  name='test2015',
                  images='test2015',
                  annotations='image_info_test2015',
                  annotation_type=AnnotationType.NONE,
              ),
          ],
      ),
      CocoConfig(
          name='2017',
          description=DESCRIPTION.format(year=2017),
          splits=[
              Split(
                  name=tfds.Split.TRAIN,
                  images='train2017',
                  annotations='annotations_trainval2017',
                  annotation_type=AnnotationType.BBOXES,
              ),
              Split(
                  name=tfds.Split.VALIDATION,
                  images='val2017',
                  annotations='annotations_trainval2017',
                  annotation_type=AnnotationType.BBOXES,
              ),
              Split(
                  name=tfds.Split.TEST,
                  images='test2017',
                  annotations='image_info_test2017',
                  annotation_type=AnnotationType.NONE,
              ),
          ],
      ),
      CocoConfig(
          name='2017_panoptic',
          description=DESCRIPTION.format(year=2017),
          has_panoptic=True,
          splits=[
              Split(
                  name=tfds.Split.TRAIN,
                  images='train2017',
                  annotations='panoptic_annotations_trainval2017',
                  annotation_type=AnnotationType.PANOPTIC,
              ),
              Split(
                  name=tfds.Split.VALIDATION,
                  images='val2017',
                  annotations='panoptic_annotations_trainval2017',
                  annotation_type=AnnotationType.PANOPTIC,
              ),
          ],
      ),
  ]

  def _info(self):

    features = {
        # Images can have variable shape
        'image': tfds.features.Image(encoding_format='jpeg'),
        'image/filename': tfds.features.Text(),
        'image/id': tf.int64,
    }
    # Either uses panotptic or original annotations
    if self.builder_config.has_panoptic:
      features.update({
          'panoptic_image': tfds.features.Image(encoding_format='png'),
          'panoptic_image/filename': tfds.features.Text(),
          'panoptic_objects': tfds.features.Sequence({
              'id': tf.int64,
              # Coco has unique id for each annotation. The id can be used for
              # mapping panoptic image to semantic segmentation label.
              'area': tf.int64,
              'bbox': tfds.features.BBoxFeature(),
              # Coco2017 has 200 categories but only 133 are present in the
              # dataset
              'label': tfds.features.ClassLabel(num_classes=133),
              'is_crowd': tf.bool,
          }),
      })
    else:
      features.update({
          'objects': tfds.features.Sequence({
              'id': tf.int64,
              # Coco has unique id for each annotation. The id can be used for
              # mapping panoptic image to semantic segmentation label.
              'area': tf.int64,
              'bbox': tfds.features.BBoxFeature(),
              # Coco has 91 categories but only 80 are present in the dataset
              'label': tfds.features.ClassLabel(num_classes=80),
              'is_crowd': tf.bool,
          }),
      })
    # More info could be added, like segmentation (as png mask), captions,
    # person key-points, more metadata (original flickr url,...).

    return tfds.core.DatasetInfo(
        builder=self,
        description=self.builder_config.description,
        # More info could be added, like the segmentation (as png mask),
        # captions, person key-points. For caption encoding, it would probably
        # be better to have a separate class CocoCaption2014 to avoid poluting
        # the main class with builder config for each encoder.
        features=tfds.features.FeaturesDict(features),
        homepage='http://cocodataset.org/#home',
        citation=_CITATION,
    )

  def _split_generators(self, dl_manager):
    """Returns SplitGenerators."""

    # Merge urls from all splits together
    urls = {}
    for split in self.builder_config.splits:
      urls['{}_images'.format(split.name)] = 'zips/{}.zip'.format(split.images)
      urls['{}_annotations'.format(split.name)] = 'annotations/{}.zip'.format(
          split.annotations)

    # DownloadManager memoize the url, so duplicate urls will only be downloaded
    # once.
    root_url = 'http://images.cocodataset.org/'
    extracted_paths = dl_manager.download_and_extract({
        key: root_url + url for key, url in urls.items()
    })

    splits = []
    for split in self.builder_config.splits:
      image_dir = extracted_paths['{}_images'.format(split.name)]
      annotations_dir = extracted_paths['{}_annotations'.format(split.name)]
      if self.builder_config.has_panoptic:
        panoptic_image_zip_path = os.path.join(
            annotations_dir,
            'annotations',
            'panoptic_{}.zip'.format(split.images)
        )
        panoptic_dir = dl_manager.extract(panoptic_image_zip_path)
        panoptic_dir = os.path.join(
            panoptic_dir, 'panoptic_{}'.format(split.images))
      else:
        panoptic_dir = None
      splits.append(tfds.core.SplitGenerator(
          name=split.name,
          gen_kwargs=dict(
              image_dir=image_dir,
              annotation_dir=annotations_dir,
              split_name=split.images,
              annotation_type=split.annotation_type,
              panoptic_dir=panoptic_dir,
          ),
      ))
    return splits

  def _generate_examples(
      self,
      image_dir,
      annotation_dir,
      split_name,
      annotation_type,
      panoptic_dir):
    """Generate examples as dicts.

    Args:
      image_dir: `str`, directory containing the images
      annotation_dir: `str`, directory containing annotations
      split_name: `str`, <split_name><year> (ex: train2014, val2017)
      annotation_type: `AnnotationType`, the annotation format (NONE, BBOXES,
        PANOPTIC)
      panoptic_dir: If annotation_type is PANOPTIC, contains the panoptic
        image directory

    Yields:
      example key and data
    """

    if annotation_type == AnnotationType.BBOXES:
      instance_filename = 'instances_{}.json'
    elif annotation_type == AnnotationType.PANOPTIC:
      instance_filename = 'panoptic_{}.json'
    elif annotation_type == AnnotationType.NONE:  # No annotation for test sets
      instance_filename = 'image_info_{}.json'

    # Load the annotations (label names, images metadata,...)
    instance_path = os.path.join(
        annotation_dir,
        'annotations',
        instance_filename.format(split_name),
    )
    coco_annotation = ANNOTATION_CLS[annotation_type](instance_path)
    # Each category is a dict:
    # {
    #    'id': 51,  # From 1-91, some entry missing
    #    'name': 'bowl',
    #    'supercategory': 'kitchen',
    # }
    categories = coco_annotation.categories
    # Each image is a dict:
    # {
    #     'id': 262145,
    #     'file_name': 'COCO_train2017_000000262145.jpg'
    #     'flickr_url': 'http://farm8.staticflickr.com/7187/xyz.jpg',
    #     'coco_url': 'http://images.cocodataset.org/train2017/xyz.jpg',
    #     'license': 2,
    #     'date_captured': '2013-11-20 02:07:55',
    #     'height': 427,
    #     'width': 640,
    # }
    images = coco_annotation.images

    # TODO(b/121375022): ClassLabel names should also contains 'id' and
    # and 'supercategory' (in addition to 'name')
    # Warning: As Coco only use 80 out of the 91 labels, the c['id'] and
    # dataset names ids won't match.
    if self.builder_config.has_panoptic:
      objects_key = 'panoptic_objects'
    else:
      objects_key = 'objects'
    self.info.features[objects_key]['label'].names = [
        c['name'] for c in categories
    ]
    # TODO(b/121375022): Conversion should be done by ClassLabel
    categories_id2name = {c['id']: c['name'] for c in categories}

    # Iterate over all images
    annotation_skipped = 0
    for image_info in sorted(images, key=lambda x: x['id']):
      if annotation_type == AnnotationType.BBOXES:
        # Each instance annotation is a dict:
        # {
        #     'iscrowd': 0,
        #     'bbox': [116.95, 305.86, 285.3, 266.03],
        #     'image_id': 480023,
        #     'segmentation': [[312.29, 562.89, 402.25, ...]],
        #     'category_id': 58,
        #     'area': 54652.9556,
        #     'id': 86,
        # }
        instances = coco_annotation.get_annotations(img_id=image_info['id'])
      elif annotation_type == AnnotationType.PANOPTIC:
        # Each panoptic annotation is a dict:
        # {
        #     'file_name': '000000037777.png',
        #     'image_id': 37777,
        #     'segments_info': [
        #         {
        #             'area': 353,
        #             'category_id': 52,
        #             'iscrowd': 0,
        #             'id': 6202563,
        #             'bbox': [221, 179, 37, 27],
        #         },
        #         ...
        #     ]
        # }
        panoptic_annotation = coco_annotation.get_annotations(
            img_id=image_info['id'])
        instances = panoptic_annotation['segments_info']
      else:
        instances = []  # No annotations

      if not instances:
        annotation_skipped += 1

      def build_bbox(x, y, width, height):
        # pylint: disable=cell-var-from-loop
        # build_bbox is only used within the loop so it is ok to use image_info
        return tfds.features.BBox(
            ymin=y / image_info['height'],
            xmin=x / image_info['width'],
            ymax=(y + height) / image_info['height'],
            xmax=(x + width) / image_info['width'],
        )
        # pylint: enable=cell-var-from-loop

      example = {
          'image': os.path.join(image_dir, split_name, image_info['file_name']),
          'image/filename': image_info['file_name'],
          'image/id': image_info['id'],
          objects_key: [{   # pylint: disable=g-complex-comprehension
              'id': instance['id'],
              'area': instance['area'],
              'bbox': build_bbox(*instance['bbox']),
              'label': categories_id2name[instance['category_id']],
              'is_crowd': bool(instance['iscrowd']),
          } for instance in instances]
      }
      if self.builder_config.has_panoptic:
        panoptic_filename = panoptic_annotation['file_name']
        panoptic_image_path = os.path.join(panoptic_dir, panoptic_filename)
        example['panoptic_image'] = panoptic_image_path
        example['panoptic_image/filename'] = panoptic_filename

      yield image_info['file_name'], example

    logging.info(
        '%d/%d images do not contains any annotations',
        annotation_skipped,
        len(images),
    )


class CocoAnnotation(object):
  """Coco annotation helper class."""

  def __init__(self, annotation_path):
    with tf.io.gfile.GFile(annotation_path) as f:
      data = json.load(f)
    self._data = data

  @property
  def categories(self):
    """Return the category dicts, as sorted in the file."""
    return self._data['categories']

  @property
  def images(self):
    """Return the image dicts, as sorted in the file."""
    return self._data['images']

  def get_annotations(self, img_id):
    """Return all annotations associated with the image id string."""
    raise NotImplementedError  # AnotationType.NONE don't have annotations


class CocoAnnotationBBoxes(CocoAnnotation):
  """Coco annotation helper class."""

  def __init__(self, annotation_path):
    super(CocoAnnotationBBoxes, self).__init__(annotation_path)

    img_id2annotations = collections.defaultdict(list)
    for a in self._data['annotations']:
      img_id2annotations[a['image_id']].append(a)
    self._img_id2annotations = {
        k: list(sorted(v, key=lambda a: a['id']))
        for k, v in img_id2annotations.items()
    }

  def get_annotations(self, img_id):
    """Return all annotations associated with the image id string."""
    # Some images don't have any annotations. Return empty list instead.
    return self._img_id2annotations.get(img_id, [])


class CocoAnnotationPanoptic(CocoAnnotation):
  """Coco annotation helper class."""

  def __init__(self, annotation_path):
    super(CocoAnnotationPanoptic, self).__init__(annotation_path)
    self._img_id2annotations = {
        a['image_id']: a for a in self._data['annotations']
    }

  def get_annotations(self, img_id):
    """Return all annotations associated with the image id string."""
    return self._img_id2annotations[img_id]


ANNOTATION_CLS = {
    AnnotationType.NONE: CocoAnnotation,
    AnnotationType.BBOXES: CocoAnnotationBBoxes,
    AnnotationType.PANOPTIC: CocoAnnotationPanoptic,
}