# coding=utf-8
# Copyright 2019 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.

"""SUN (Scene UNderstanding) datasets."""

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

import io
import os

from absl import logging
import numpy as np
import tensorflow as tf
from tensorflow_datasets.core import utils
import tensorflow_datasets.public_api as tfds


_SUN397_CITATION = """\
@INPROCEEDINGS{Xiao:2010,
author={J. {Xiao} and J. {Hays} and K. A. {Ehinger} and A. {Oliva} and A. {Torralba}},
booktitle={2010 IEEE Computer Society Conference on Computer Vision and Pattern Recognition},
title={SUN database: Large-scale scene recognition from abbey to zoo},
year={2010},
volume={},
number={},
pages={3485-3492},
keywords={computer vision;human factors;image classification;object recognition;visual databases;SUN database;large-scale scene recognition;abbey;zoo;scene categorization;computer vision;scene understanding research;scene category;object categorization;scene understanding database;state-of-the-art algorithms;human scene classification performance;finer-grained scene representation;Sun;Large-scale systems;Layout;Humans;Image databases;Computer vision;Anthropometry;Bridges;Legged locomotion;Spatial databases}, 
doi={10.1109/CVPR.2010.5539970},
ISSN={1063-6919},
month={June},}
"""
_SUN397_DESCRIPTION = """\
The database contains 108,753 images of 397 categories, used in the
Scene UNderstanding (SUN) benchmark. The number of images varies across
categories, but there are at least 100 images per category.

Several configs of the dataset are made available through TFDS:
- A custom (random) partition of the whole dataset with 76,128 training images,
  10,875 validation images and 21,750 test images. Images have been resized to
  have at most 120,000 pixels, and encoded as JPEG with quality of 72.
- "standard-part1-120k", "standard-part2-120k", ..., "standard-part10-120k":
  Each of the 10 official train/test partitions with 50 images per class in each
  split. Images have been resized to have at most 120,000 pixels, and encoded
  as JPEG with quality of 72.
"""
_SUN397_URL = "https://vision.princeton.edu/projects/2010/SUN/"

# These images are badly encoded and cannot be decoded correctly (TF), or the
# decoding is not deterministic (PIL).
_SUN397_IGNORE_IMAGES = [
    "SUN397/c/church/outdoor/sun_bhenjvsvrtumjuri.jpg",
]

_SUN397_BUILDER_CONFIG_DESCRIPTION_PATTERN = (
    "Train and test splits from the official partition number %d. "
    "Images are resized to have at most %s pixels, and compressed with 72 JPEG "
    "quality.")


def _decode_image(fobj, session, filename):
  """Reads and decodes an image from a file object as a Numpy array.

  The SUN dataset contains images in several formats (despite the fact that
  all of them have .jpg extension). Some of them are:
    - BMP (RGB)
    - PNG (grayscale, RGBA, RGB interlaced)
    - JPEG (RGB)
    - GIF (1-frame RGB)
  Since TFDS assumes that all images have the same number of channels, we
  convert all of them to RGB.

  Args:
    fobj: File object to read from.
    session: TF session used to decode the images.
    filename: Filename of the original image in the archive.

  Returns:
    Numpy array with shape (height, width, channels).
  """

  buf = fobj.read()
  image = tfds.core.lazy_imports.cv2.imdecode(
      np.fromstring(buf, dtype=np.uint8), flags=3)  # Note: Converts to RGB.
  if image is None:
    logging.warning(
        "Image %s could not be decoded by OpenCV, falling back to TF", filename)
    try:
      image = tf.image.decode_image(buf, channels=3)
      image = session.run(image)
    except tf.errors.InvalidArgumentError:
      logging.fatal("Image %s could not be decoded by Tensorflow", filename)

  # The GIF images contain a single frame.
  if len(image.shape) == 4:  # rank=4 -> rank=3
    image = image.reshape(image.shape[1:])

  return image


def _encode_jpeg(image, quality=None):
  cv2 = tfds.core.lazy_imports.cv2
  extra_args = [[int(cv2.IMWRITE_JPEG_QUALITY), quality]] if quality else []
  _, buff = cv2.imencode(".jpg", image, *extra_args)
  return io.BytesIO(buff.tostring())


def _process_image_file(
    fobj, session, filename, quality=None, target_pixels=None):
  """Process image files from the dataset."""
  # We need to read the image files and convert them to JPEG, since some files
  # actually contain GIF, PNG or BMP data (despite having a .jpg extension) and
  # some encoding options that will make TF crash in general.
  image = _decode_image(fobj, session, filename=filename)
  # Get image height and width.
  height, width, _ = image.shape
  actual_pixels = height * width
  if target_pixels and actual_pixels > target_pixels:
    factor = np.sqrt(target_pixels / actual_pixels)
    image = tfds.core.lazy_imports.cv2.resize(
        image, dsize=None, fx=factor, fy=factor)
  return _encode_jpeg(image, quality=quality)


class Sun397Config(tfds.core.BuilderConfig):
  """BuilderConfig for Sun 397 dataset."""

  # Version history:
  # 4.0.0: Pictures resized and compressed, new configs, using S3 features.
  # 3.0.0: S3 with new hashing function (different shuffle).
  # 2.0.0: S3 (new shuffling, sharding and slicing mechanism).

  def __init__(
      self, target_pixels=None, partition=None, quality=None, **kwargs):
    self._target_pixels = target_pixels
    self._partition = partition
    self._quality = quality
    super(Sun397Config, self).__init__(**kwargs)

  @property
  def target_pixels(self):
    return self._target_pixels

  @property
  def partition(self):
    return self._partition

  @property
  def quality(self):
    return self._quality


def _generate_builder_configs():
  """Return the BuilderConfig objects for the SUN397 dataset."""
  version = tfds.core.Version("4.0.0")
  builder_configs = [
      # Images randomly split into train/valid/test splits (70%/10%/20%), and
      # images resized to have at most 120,000 pixels.
      Sun397Config(
          name="tfds",
          version=version,
          target_pixels=120000,
          description=(
              "TFDS partition with random train/validation/test splits with "
              "70%/10%/20% of the images, respectively. Images are resized to "
              "have at most 120,000 pixels, and are compressed with 72 JPEG "
              "quality.")),
  ]
  # Configs for each of the standard partitions of the dataset.
  for partition in range(1, 10 + 1):
    description = _SUN397_BUILDER_CONFIG_DESCRIPTION_PATTERN % (partition,
                                                                "120,000")
    builder_configs.append(
        Sun397Config(
            name="standard-part%d-120k" % partition,
            partition=partition,
            target_pixels=120000,
            version=version,
            description=description))
  return builder_configs


class Sun397(tfds.core.GeneratorBasedBuilder):
  """Sun397 Scene Recognition Benchmark."""

  BUILDER_CONFIGS = _generate_builder_configs()

  def __init__(self, tfds_split_files=None, **kwargs):
    super(Sun397, self).__init__(**kwargs)
    # Note: Only used for tests, since the data is fake.
    if not tfds_split_files:
      tfds_split_files = {
          "tr": "sun397_tfds_tr.txt",
          "te": "sun397_tfds_te.txt",
          "va": "sun397_tfds_va.txt",
      }
      for split, filename in tfds_split_files.items():
        tfds_split_files[split] = tfds.core.get_tfds_path(
            os.path.join("image", filename))
    self._tfds_split_files = tfds_split_files

  def _info(self):
    names_file = tfds.core.get_tfds_path(
        os.path.join("image", "sun397_labels.txt"))
    return tfds.core.DatasetInfo(
        builder=self,
        description=_SUN397_DESCRIPTION,
        features=tfds.features.FeaturesDict({
            "file_name": tfds.features.Text(),
            "image": tfds.features.Image(shape=(None, None, 3)),
            "label": tfds.features.ClassLabel(names_file=names_file),
        }),
        urls=[_SUN397_URL],
        citation=_SUN397_CITATION)

  def _split_generators(self, dl_manager):
    paths = dl_manager.download_and_extract({
        "images": tfds.download.Resource(
            url=_SUN397_URL + "SUN397.tar.gz",
            extract_method=tfds.download.ExtractMethod.NO_EXTRACT),
        "partitions": _SUN397_URL + "download/Partitions.zip",
    })
    if not isinstance(paths, dict):
      # While testing download_and_extract() returns the dir containing the
      # test files.
      paths = {
          "images": os.path.join(paths, "SUN397.tar.gz"),
          "partitions": os.path.join(paths, "Partitions"),
      }
    images = tfds.download.Resource(
        path=paths["images"],
        extract_method=tfds.download.ExtractMethod.TAR_GZ_STREAM)
    if self.builder_config.name == "tfds":
      subset_images = self._get_tfds_subsets_images()
      return [
          tfds.core.SplitGenerator(
              name=tfds.Split.TRAIN,
              gen_kwargs=dict(
                  archive=dl_manager.iter_archive(images),
                  subset_images=subset_images["tr"])),
          tfds.core.SplitGenerator(
              name=tfds.Split.TEST,
              gen_kwargs=dict(
                  archive=dl_manager.iter_archive(images),
                  subset_images=subset_images["te"])),
          tfds.core.SplitGenerator(
              name=tfds.Split.VALIDATION,
              gen_kwargs=dict(
                  archive=dl_manager.iter_archive(images),
                  subset_images=subset_images["va"])),
      ]
    else:
      subset_images = self._get_partition_subsets_images(paths["partitions"])
      return [
          tfds.core.SplitGenerator(
              name=tfds.Split.TRAIN,
              gen_kwargs=dict(
                  archive=dl_manager.iter_archive(images),
                  subset_images=subset_images["tr"])),
          tfds.core.SplitGenerator(
              name=tfds.Split.TEST,
              gen_kwargs=dict(
                  archive=dl_manager.iter_archive(images),
                  subset_images=subset_images["te"])),
      ]

  def _generate_examples(self, archive, subset_images):
    prefix_len = len("SUN397")
    with tf.Graph().as_default():
      with utils.nogpu_session() as sess:
        for filepath, fobj in archive:
          # Note: all files in the tar.gz are in SUN397/...
          filename = filepath[prefix_len:]
          if filename in subset_images:
            # Example:
            # From filename: /c/car_interior/backseat/sun_aenygxwhhmjtisnf.jpg
            # To class: /c/car_interior/backseat
            label = "/".join(filename.split("/")[:-1])
            image = _process_image_file(
                fobj, sess, filename,
                quality=self.builder_config.quality,
                target_pixels=self.builder_config.target_pixels)
            record = {
                "file_name": filename,
                "image": image,
                "label": label,
            }
            yield filename, record

  def _get_tfds_subsets_images(self):
    splits_sets = {}
    for split, filepath in self._tfds_split_files.items():
      splits_sets[split] = self._load_image_set_from_file(filepath)
    return splits_sets

  def _get_partition_subsets_images(self, partitions_dir):
    # Get the ID of all images in the dataset.
    all_images = set()
    for split_images in self._get_tfds_subsets_images().values():
      all_images.update(split_images)
    # Load the images in the training/test split of this partition.
    filenames = {
        "tr": "Training_%02d.txt" % self.builder_config.partition,
        "te": "Testing_%02d.txt" % self.builder_config.partition
    }
    splits_sets = {}
    for split, filename in filenames.items():
      filepath = os.path.join(partitions_dir, filename)
      splits_sets[split] = self._load_image_set_from_file(filepath)
    # Put the remaining images in the dataset into the "validation" split.
    splits_sets["va"] = all_images - (splits_sets["tr"] | splits_sets["te"])
    return splits_sets

  def _load_image_set_from_file(self, filepath):
    with tf.io.gfile.GFile(filepath, mode="r") as f:
      return set([line.strip() for line in f])
