# 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.

"""Bounding boxes feature."""

import collections

import numpy as np
import tensorflow.compat.v2 as tf

from tensorflow_datasets.core import utils
from tensorflow_datasets.core import lazy_imports_lib
from tensorflow_datasets.core.features import feature
from tensorflow_datasets.core.utils import type_utils

Json = type_utils.Json

BBox = collections.namedtuple('BBox', 'ymin, xmin, ymax, xmax')


class BBoxFeature(feature.Tensor):
  """`FeatureConnector` for a normalized bounding box.

  Note: If you have multiple bounding boxes, you may want to wrap the feature
  inside a `tfds.feature.Sequence`.

  Input:
    * `tfds.features.BBox` tuple.

  Output:
    bbox: tf.Tensor of type `tf.float32` and shape `[4,]` which contains the
      normalized coordinates of the bounding box `[ymin, xmin, ymax, xmax]`

  Example:
    * In the DatasetInfo object:

    ```
    features=features.FeatureDict({
        'bbox': features.BBox(shape=(None, 64, 64, 3)),
    })
    ```

    * During generation:

    ```
    yield {
        'input': tfds.feature.BBox(ymin=0.3, xmin=0.8, ymax=0.5, xmax=1.0),
    }
    ```
  """

  def __init__(self):
    super(BBoxFeature, self).__init__(shape=(4,), dtype=tf.float32)

  def encode_example(self, bbox):
    """See base class for details."""
    # Validate the coordinates
    for coordinate in bbox:
      if not isinstance(coordinate, float):
        raise ValueError(
            'BBox coordinates should be float. Got {}.'.format(bbox))
      if not 0.0 <= coordinate <= 1.0:
        raise ValueError(
            'BBox coordinates should be between 0 and 1. Got {}.'.format(bbox))
      if bbox.xmax < bbox.xmin or bbox.ymax < bbox.ymin:
        raise ValueError(
            'BBox coordinates should have min <= max. Got {}.'.format(bbox))

    return super(BBoxFeature, self).encode_example(
        [bbox.ymin, bbox.xmin, bbox.ymax, bbox.xmax]
    )

  def _build_thumbnail_with_bbox(self, ex: np.ndarray):
    """Returns blank image with Bboxes drawn on it."""
    PIL_Image = lazy_imports_lib.lazy_imports.PIL_Image
    PIL_ImageDraw = lazy_imports_lib.lazy_imports.PIL_ImageDraw

    SIZE = 150
    blank_img = PIL_Image.new('RGB', (SIZE, SIZE), (255, 255, 255))
    draw = PIL_ImageDraw.Draw(blank_img)
    rs = np.random.RandomState(97531)  # freeze random state

    for i in range(ex.shape[0]):
      # Rescale coordinates to match size of blank_image
      ymin, xmin, ymax, xmax = ex[i, :]*SIZE
      # Generate random rgb values for Bbox ouline
      r, g, b = list(rs.randint(0, 256, size=3))
      draw.rectangle(((xmin, ymin), (xmax, ymax)), outline=(r, g, b))
    return blank_img

  def repr_html(self, ex: np.ndarray) -> str:
    """Returns the HTML str representation of an Image with BBoxes."""
    ex = np.expand_dims(ex, axis=0)  # Convert single bounding box to single batch.
    return self._repr_html(ex)

  def repr_html_batch(self, ex: np.ndarray) -> str:
    """Returns the HTML str representation of an Image with BBoxes (Sequence)."""
    return self._repr_html(ex)

  def _repr_html(self, ex: np.ndarray) -> str:
    """Returns the HTML str representation of an Image with BBoxes."""
    img = self._build_thumbnail_with_bbox(ex)
    img_str = utils.get_base64(lambda buff: img.save(buff, format='PNG'))
    return f'<img src="data:image/png;base64,{img_str}" alt="Img" />'

  @classmethod
  def from_json_content(cls, value: Json) -> 'BBoxFeature':
    del value  # Unused
    return cls()

  def to_json_content(self) -> Json:
    return dict()

@@ -2508,753 +2508,8 @@
  )%0A%0A
-  def _build_thumbnail_with_bbox(self, ex: np.ndarray):%0A    %22%22%22Returns blank image with Bboxes drawn on it.%22%22%22%0A    PIL_Image = lazy_imports_lib.lazy_imports.PIL_Image%0A    PIL_ImageDraw = lazy_imports_lib.lazy_imports.PIL_ImageDraw%0A%0A    SIZE = 150%0A    blank_img = PIL_Image.new('RGB', (SIZE, SIZE), (255, 255, 255))%0A    draw = PIL_ImageDraw.Draw(blank_img)%0A    rs = np.random.RandomState(97531)  # freeze random state%0A%0A    for i in range(ex.shape%5B0%5D):%0A      # Rescale coordinates to match size of blank_image%0A      ymin, xmin, ymax, xmax = ex%5Bi, :%5D*SIZE%0A      # Generate random rgb values for Bbox ouline%0A      r, g, b = list(rs.randint(0, 256, size=3))%0A      draw.rectangle(((xmin, ymin), (xmax, ymax)), outline=(r, g, b))%0A    return blank_img%0A%0A
   de
@@ -2542,32 +2542,32 @@
 darray) -%3E str:%0A
+
     %22%22%22Returns t
@@ -3031,21 +3031,16 @@
   img = 
-self.
 _build_t
@@ -3361,16 +3361,16 @@
 %3E Json:%0A
-
     retu
@@ -3379,8 +3379,717 @@
  dict()%0A
+%0Adef _build_thumbnail_with_bbox(ex: np.ndarray):%0A  %22%22%22Returns blank image with Bboxes drawn on it.%22%22%22%0A  PIL_Image = lazy_imports_lib.lazy_imports.PIL_Image%0A  PIL_ImageDraw = lazy_imports_lib.lazy_imports.PIL_ImageDraw%0A%0A  SIZE = 150%0A  blank_img = PIL_Image.new('RGB', (SIZE, SIZE), (255, 255, 255))%0A  draw = PIL_ImageDraw.Draw(blank_img)%0A  rs = np.random.RandomState(97531)  # freeze random state%0A%0A  for i in range(ex.shape%5B0%5D):%0A    # Rescale coordinates to match size of blank_image%0A    ymin, xmin, ymax, xmax = ex%5Bi, :%5D*SIZE%0A    # Generate random rgb values for Bbox ouline%0A    r, g, b = list(rs.randint(0, 256, size=3))%0A    draw.rectangle(((xmin, ymin), (xmax, ymax)), outline=(r, g, b))%0A  return blank_img%0A

Make `_build_thumbnail_with_bbox` a function