Dataset modules (#9903)

* Dataset utilities added.

* Global model definition

* Dataset modules added.

* Dataset modules fix.

* global features model training added

* global features fix

* Test dataset update

* PR fixes

* repo sync

* repo sync

* Syncing 2

* Syncing 2

* Global model definition added

* Global model definition added, synced

* Adding global model dataset related modules

* Global model training

* tensorboard module added

* linting issues fixed

* linting fixes.

* linting fixes.

* Fix for previous PR

* PR fixes

* Minor fixes

* Minor fixes

* Dataset download fix

* Comments fix

* sfm120k fix

* sfm120k fix

* names fix

* Update

* Update

* Merge branch 'global_model_training'

# Conflicts:
#	research/delf/delf/python/datasets/generic_dataset.py
#	research/delf/delf/python/datasets/generic_dataset_test.py
#	research/delf/delf/python/datasets/sfm120k/__init__.py
#	research/delf/delf/python/datasets/sfm120k/sfm120k.py
#	research/delf/delf/python/datasets/sfm120k/sfm120k_test.py
#	research/delf/delf/python/datasets/tuples_dataset.py
#	research/delf/delf/python/datasets/tuples_dataset_test.py
#	research/delf/delf/python/training/global_features/__init__.py
#	research/delf/delf/python/training/global_features/train.py
#	research/delf/delf/python/training/model/global_model.py
#	research/delf/delf/python/training/model/global_model_test.py
#	research/delf/delf/python/training/tensorboard_utils.py

* Merge branch 'global_model_training'

# Conflicts:
#	research/delf/delf/python/datasets/generic_dataset.py
#	research/delf/delf/python/datasets/generic_dataset_test.py
#	research/delf/delf/python/datasets/sfm120k/__init__.py
#	research/delf/delf/python/datasets/sfm120k/sfm120k.py
#	research/delf/delf/python/datasets/sfm120k/sfm120k_test.py
#	research/delf/delf/python/datasets/tuples_dataset.py
#	research/delf/delf/python/datasets/tuples_dataset_test.py
#	research/delf/delf/python/training/global_features/__init__.py
#	research/delf/delf/python/training/global_features/train.py
#	research/delf/delf/python/training/model/global_model.py
#	research/delf/delf/python/training/model/global_model_test.py
#	research/delf/delf/python/training/tensorboard_utils.py

* PR fixes

* Merge branch 'global_model_training'

# Conflicts:
#	research/delf/delf/python/datasets/generic_dataset.py
#	research/delf/delf/python/datasets/generic_dataset_test.py
#	research/delf/delf/python/datasets/sfm120k/__init__.py
#	research/delf/delf/python/datasets/sfm120k/sfm120k.py
#	research/delf/delf/python/datasets/sfm120k/sfm120k_test.py
#	research/delf/delf/python/datasets/tuples_dataset.py
#	research/delf/delf/python/datasets/tuples_dataset_test.py
#	research/delf/delf/python/training/global_features/__init__.py
#	research/delf/delf/python/training/global_features/train.py
#	research/delf/delf/python/training/model/global_model.py
#	research/delf/delf/python/training/model/global_model_test.py
#	research/delf/delf/python/training/tensorboard_utils.py

* Merge branch 'global_model_training'

# Conflicts:
#	research/delf/delf/python/datasets/generic_dataset.py
#	research/delf/delf/python/datasets/generic_dataset_test.py
#	research/delf/delf/python/datasets/sfm120k/__init__.py
#	research/delf/delf/python/datasets/sfm120k/sfm120k.py
#	research/delf/delf/python/datasets/sfm120k/sfm120k_test.py
#	research/delf/delf/python/datasets/tuples_dataset.py
#	research/delf/delf/python/datasets/tuples_dataset_test.py
#	research/delf/delf/python/training/global_features/__init__.py
#	research/delf/delf/python/training/global_features/train.py
#	research/delf/delf/python/training/model/global_model.py
#	research/delf/delf/python/training/model/global_model_test.py
#	research/delf/delf/python/training/tensorboard_utils.py

* Merge branch 'global_model_training'

# Conflicts:
#	research/delf/delf/python/datasets/generic_dataset.py
#	research/delf/delf/python/datasets/generic_dataset_test.py
#	research/delf/delf/python/datasets/sfm120k/__init__.py
#	research/delf/delf/python/datasets/sfm120k/sfm120k.py
#	research/delf/delf/python/datasets/sfm120k/sfm120k_test.py
#	research/delf/delf/python/datasets/tuples_dataset.py
#	research/delf/delf/python/datasets/tuples_dataset_test.py
#	research/delf/delf/python/training/global_features/__init__.py
#	research/delf/delf/python/training/global_features/train.py
#	research/delf/delf/python/training/model/global_model.py
#	research/delf/delf/python/training/model/global_model_test.py
#	research/delf/delf/python/training/tensorboard_utils.py

research/delf/delf/python/datasets/generic_dataset.py

+# Lint as: python3
+# Copyright 2021 The TensorFlow 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.
+# ==============================================================================
+"""Functions for generic image dataset creation."""
+
+import os
+
+from delf.python.datasets import utils
+
+
+class ImagesFromList():
+  """A generic data loader that loads images from a list.
+
+  Supports images of different sizes.
+  """
+
+  def __init__(self, root, image_paths, imsize=None, bounding_boxes=None,
+               loader=utils.default_loader):
+    """ImagesFromList object initialization.
+
+    Args:
+      root: String, root directory path.
+      image_paths: List, relative image paths as strings.
+      imsize: Integer, defines the maximum size of longer image side.
+      bounding_boxes: List of (x1,y1,x2,y2) tuples to crop the query images.
+      loader: Callable, a function to load an image given its path.
+
+    Raises:
+      ValueError: Raised if `image_paths` list is empty.
+    """
+    # List of the full image filenames.
+    images_filenames = [os.path.join(root, image_path) for image_path in
+                        image_paths]
+
+    if not images_filenames:
+      raise ValueError("Dataset contains 0 images.")
+
+    self.root = root
+    self.images = image_paths
+    self.imsize = imsize
+    self.images_filenames = images_filenames
+    self.bounding_boxes = bounding_boxes
+    self.loader = loader
+
+  def __getitem__(self, index):
+    """Called to load an image at the given `index`.
+
+    Args:
+        index: Integer, image index.
+
+    Returns:
+        image: Tensor, loaded image.
+    """
+    path = self.images_filenames[index]
+
+    if self.bounding_boxes is not None:
+      img = self.loader(path, self.imsize, self.bounding_boxes[index])
+    else:
+      img = self.loader(path, self.imsize)
+
+    return img
+
+  def __len__(self):
+    """Implements the built-in function len().
+
+    Returns:
+      len: Number of images in the dataset.
+    """
+    return len(self.images_filenames)

research/delf/delf/python/datasets/generic_dataset_test.py

+# Lint as: python3
+# Copyright 2021 The TensorFlow 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.
+# ==============================================================================
+"""Tests for generic dataset."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from absl import flags
+
+import numpy as np
+from PIL import Image
+import tensorflow as tf
+
+from delf.python.datasets import generic_dataset
+
+FLAGS = flags.FLAGS
+
+
+class GenericDatasetTest(tf.test.TestCase):
+  """Test functions for generic dataset."""
+
+  def testGenericDataset(self):
+    """Tests loading dummy images from list."""
+    # Number of images to be created.
+    n = 2
+    image_names = []
+
+    # Create and save `n` dummy images.
+    for i in range(n):
+      dummy_image = np.random.rand(1024, 750, 3) * 255
+      img_out = Image.fromarray(dummy_image.astype('uint8')).convert('RGB')
+      filename = os.path.join(FLAGS.test_tmpdir,
+                              'test_image_{}.jpg'.format(i))
+      img_out.save(filename)
+      image_names.append('test_image_{}.jpg'.format(i))
+
+    data = generic_dataset.ImagesFromList(root=FLAGS.test_tmpdir,
+                                          image_paths=image_names,
+                                          imsize=1024)
+    self.assertLen(data, n)
+
+
+if __name__ == '__main__':
+  tf.test.main()

research/delf/delf/python/datasets/sfm120k/__init__.py

+# Copyright 2021 The TensorFlow 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.
+# ==============================================================================
+"""Module exposing Sfm120k dataset for training."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# pylint: disable=unused-import
+from delf.python.datasets.sfm120k import sfm120k
+# pylint: enable=unused-import

research/delf/delf/python/datasets/sfm120k/dataset_download.py

+# Copyright 2021 The TensorFlow 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.
+# ==============================================================================
+"""Structure-from-Motion dataset (Sfm120k) download function."""
+
+import os
+
+import tensorflow as tf
+
+
+def download_train(data_dir):
+  """Checks, and, if required, downloads the necessary files for the training.
+
+  Checks if the data necessary for running the example training script exist.
+  If not, it downloads it in the following folder structure:
+    DATA_ROOT/train/retrieval-SfM-120k/ : folder with rsfm120k images and db
+      files.
+    DATA_ROOT/train/retrieval-SfM-30k/  : folder with rsfm30k images and db
+      files.
+  """
+
+  # Create data folder if does not exist.
+  if not tf.io.gfile.exists(data_dir):
+    tf.io.gfile.mkdir(data_dir)
+
+  # Create datasets folder if does not exist.
+  datasets_dir = os.path.join(data_dir, 'train')
+  if not tf.io.gfile.exists(datasets_dir):
+    tf.io.gfile.mkdir(datasets_dir)
+
+  # Download folder train/retrieval-SfM-120k/.
+  src_dir = 'http://cmp.felk.cvut.cz/cnnimageretrieval/data/train/ims'
+  dst_dir = os.path.join(datasets_dir, 'retrieval-SfM-120k', 'ims')
+  download_file = 'ims.tar.gz'
+  if not tf.io.gfile.exists(dst_dir):
+    src_file = os.path.join(src_dir, download_file)
+    dst_file = os.path.join(dst_dir, download_file)
+    print('>> Image directory does not exist. Creating: {}'.format(dst_dir))
+    tf.io.gfile.makedirs(dst_dir)
+    print('>> Downloading ims.tar.gz...')
+    os.system('wget {} -O {}'.format(src_file, dst_file))
+    print('>> Extracting {}...'.format(dst_file))
+    os.system('tar -zxf {} -C {}'.format(dst_file, dst_dir))
+    print('>> Extracted, deleting {}...'.format(dst_file))
+    os.system('rm {}'.format(dst_file))
+
+  # Create symlink for train/retrieval-SfM-30k/.
+  dst_dir_old = os.path.join(datasets_dir, 'retrieval-SfM-120k', 'ims')
+  dst_dir = os.path.join(datasets_dir, 'retrieval-SfM-30k', 'ims')
+  if not (tf.io.gfile.exists(dst_dir) or os.path.islink(dst_dir)):
+    tf.io.gfile.makedirs(os.path.join(datasets_dir, 'retrieval-SfM-30k'))
+    os.system('ln -s {} {}'.format(dst_dir_old, dst_dir))
+    print(
+            '>> Created symbolic link from retrieval-SfM-120k/ims to '
+            'retrieval-SfM-30k/ims')
+
+  # Download db files.
+  src_dir = 'http://cmp.felk.cvut.cz/cnnimageretrieval/data/train/dbs'
+  datasets = ['retrieval-SfM-120k', 'retrieval-SfM-30k']
+  for dataset in datasets:
+    dst_dir = os.path.join(datasets_dir, dataset)
+    if dataset == 'retrieval-SfM-120k':
+      download_files = ['{}.pkl'.format(dataset),
+                        '{}-whiten.pkl'.format(dataset)]
+      download_eccv2020 = '{}-val-eccv2020.pkl'.format(dataset)
+    elif dataset == 'retrieval-SfM-30k':
+      download_files = ['{}-whiten.pkl'.format(dataset)]
+      download_eccv2020 = None
+
+    if not tf.io.gfile.exists(dst_dir):
+      print('>> Dataset directory does not exist. Creating: {}'.format(
+              dst_dir))
+      tf.io.gfile.mkdir(dst_dir)
+
+    for i in range(len(download_files)):
+      src_file = os.path.join(src_dir, download_files[i])
+      dst_file = os.path.join(dst_dir, download_files[i])
+      if not os.path.isfile(dst_file):
+        print('>> DB file {} does not exist. Downloading...'.format(
+                download_files[i]))
+        os.system('wget {} -O {}'.format(src_file, dst_file))
+
+      if download_eccv2020:
+        eccv2020_dst_file = os.path.join(dst_dir, download_eccv2020)
+        if not os.path.isfile(eccv2020_dst_file):
+          eccv2020_src_dir = \
+            "http://ptak.felk.cvut.cz/personal/toliageo/share/how/dataset/"
+          eccv2020_dst_file = os.path.join(dst_dir, download_eccv2020)
+          eccv2020_src_file = os.path.join(eccv2020_src_dir,
+                                           download_eccv2020)
+          os.system('wget {} -O {}'.format(eccv2020_src_file,
+                                           eccv2020_dst_file))

research/delf/delf/python/datasets/sfm120k/sfm120k.py

+# Copyright 2021 The TensorFlow 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.
+# ==============================================================================
+"""Structure-from-Motion dataset (Sfm120k) module.
+
+[1] From Single Image Query to Detailed 3D Reconstruction.
+Johannes L. Schonberger, Filip Radenovic, Ondrej Chum, Jan-Michael Frahm.
+The related paper can be found at: https://ieeexplore.ieee.org/document/7299148.
+"""
+
+import os
+import pickle
+import tensorflow as tf
+
+from delf.python.datasets import tuples_dataset
+from delf.python.datasets import utils
+
+
+def id2filename(image_id, prefix):
+  """Creates a training image path out of its id name.
+
+  Used for the image mapping in the Sfm120k datset.
+
+  Args:
+    image_id: String, image id.
+    prefix: String, root directory where images are saved.
+
+  Returns:
+   filename: String, full image filename.
+  """
+  if prefix:
+    return os.path.join(prefix, image_id[-2:], image_id[-4:-2], image_id[-6:-4],
+                        image_id)
+  else:
+    return os.path.join(image_id[-2:], image_id[-4:-2], image_id[-6:-4],
+                        image_id)
+
+
+class _Sfm120k(tuples_dataset.TuplesDataset):
+  """Structure-from-Motion (Sfm120k) dataset instance.
+
+  The dataset contains the image names lists for training and validation,
+  the cluster ID (3D model ID) for each image and indices forming
+  query-positive pairs of images. The images are loaded per epoch and resized
+  on the fly to the desired dimensionality.
+  """
+
+  def __init__(self, mode, data_root, imsize=None, num_negatives=5,
+               num_queries=2000, pool_size=20000, loader=utils.default_loader,
+               eccv2020=False):
+    """Structure-from-Motion (Sfm120k) dataset initialization.
+
+    Args:
+      mode: Either 'train' or 'val'.
+      data_root: Path to the root directory of the dataset.
+      imsize: Integer, defines the maximum size of longer image side.
+      num_negatives: Integer, number of negative images per one query.
+      num_queries: Integer, number of query images.
+      pool_size: Integer, size of the negative image pool, from where the
+        hard-negative images are chosen.
+      loader: Callable, a function to load an image given its path.
+      eccv2020: Bool, whether to use a new validation dataset used with ECCV
+        2020 paper (https://arxiv.org/abs/2007.13172).
+
+    Raises:
+      ValueError: Raised if `mode` is not one of 'train' or 'val'.
+    """
+    if mode not in ['train', 'val']:
+      raise ValueError(
+              "`mode` argument should be either 'train' or 'val', passed as a "
+              "String.")
+
+    # Setting up the paths for the dataset.
+    if eccv2020:
+      name = "retrieval-SfM-120k-val-eccv2020"
+    else:
+      name = "retrieval-SfM-120k"
+    db_root = os.path.join(data_root, 'train/retrieval-SfM-120k')
+    ims_root = os.path.join(db_root, 'ims/')
+
+    # Loading the dataset db file.
+    db_filename = os.path.join(db_root, '{}.pkl'.format(name))
+
+    with tf.io.gfile.GFile(db_filename, 'rb') as f:
+      db = pickle.load(f)[mode]
+
+    # Setting full paths for the dataset images.
+    self.images = [id2filename(img_name, None) for
+                   img_name in db['cids']]
+
+    # Initializing tuples dataset.
+    super().__init__(name, mode, db_root, imsize, num_negatives, num_queries,
+                     pool_size, loader, ims_root)
+
+  def Sfm120kInfo(self):
+    """Metadata for the Sfm120k dataset.
+
+    The dataset contains the image names lists for training and
+    validation, the cluster ID (3D model ID) for each image and indices
+    forming query-positive pairs of images. The images are loaded per epoch
+    and resized on the fly to the desired dimensionality.
+
+    Returns:
+      info: dictionary with the dataset parameters.
+    """
+    info = {'train': {'clusters': 91642, 'pidxs': 181697, 'qidxs': 181697},
+            'val': {'clusters': 6403, 'pidxs': 1691, 'qidxs': 1691}}
+    return info
+
+
+def CreateDataset(mode, data_root, imsize=None, num_negatives=5,
+                  num_queries=2000, pool_size=20000,
+                  loader=utils.default_loader, eccv2020=False):
+  '''Creates Structure-from-Motion (Sfm120k) dataset.
+
+  Args:
+    mode: String, either 'train' or 'val'.
+    data_root: Path to the root directory of the dataset.
+    imsize: Integer, defines the maximum size of longer image side.
+    num_negatives: Integer, number of negative images per one query.
+    num_queries: Integer, number of query images.
+    pool_size: Integer, size of the negative image pool, from where the
+      hard-negative images are chosen.
+    loader: Callable, a function to load an image given its path.
+    eccv2020: Bool, whether to use a new validation dataset used with ECCV
+      2020 paper (https://arxiv.org/abs/2007.13172).
+
+  Returns:
+    sfm120k: Sfm120k dataset instance.
+  '''
+  return _Sfm120k(mode, data_root, imsize, num_negatives, num_queries,
+                  pool_size, loader, eccv2020)

research/delf/delf/python/datasets/sfm120k/sfm120k_test.py

+# Lint as: python3
+# Copyright 2021 The TensorFlow 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.
+# ==============================================================================
+"""Tests for Sfm120k dataset module."""
+
+import tensorflow as tf
+
+from delf.python.datasets.sfm120k import sfm120k
+
+
+class Sfm120kTest(tf.test.TestCase):
+  """Tests for Sfm120k dataset module."""
+
+  def testId2Filename(self):
+    """Tests conversion of image id to full path mapping."""
+    image_id = "29fdc243aeb939388cfdf2d081dc080e"
+    prefix = "train/retrieval-SfM-120k/ims/"
+    path = sfm120k.id2filename(image_id, prefix)
+    expected_path = "train/retrieval-SfM-120k/ims/0e/08/dc" \
+                    "/29fdc243aeb939388cfdf2d081dc080e"
+    self.assertEqual(path, expected_path)
+
+
+if __name__ == '__main__':
+  tf.test.main()

research/delf/delf/python/datasets/tuples_dataset.py

+# Lint as: python3
+# Copyright 2021 The TensorFlow 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.
+# ==============================================================================
+"""Tuple dataset module.
+
+Based on the Radenovic et al. ECCV16: CNN image retrieval learns from BoW.
+For more information refer to https://arxiv.org/abs/1604.02426.
+"""
+
+import os
+import pickle
+
+import numpy as np
+import tensorflow as tf
+
+from delf.python.datasets import utils as image_loading_utils
+from delf.python.training import global_features_utils
+from delf.python.training.model import global_model
+
+
+class TuplesDataset():
+  """Data loader that loads training and validation tuples.
+
+  After initialization, the function create_epoch_tuples() should be called to
+  create the dataset tuples. After that, the dataset can be iterated through
+  using next() function.
+  Tuples are based on Radenovic et al. ECCV16 work: CNN image retrieval
+  learns from BoW. For more information refer to
+  https://arxiv.org/abs/1604.02426.
+  """
+
+  def __init__(self, name, mode, data_root, imsize=None, num_negatives=5,
+               num_queries=2000, pool_size=20000,
+               loader=image_loading_utils.default_loader, ims_root=None):
+    """TuplesDataset object initialization.
+
+    Args:
+      name: String, dataset name. I.e. 'retrieval-sfm-120k'.
+      mode: 'train' or 'val' for training and validation parts of dataset.
+      data_root: Path to the root directory of the dataset.
+      imsize: Integer, defines the maximum size of longer image side transform.
+      num_negatives: Integer, number of negative images for a query image in a
+        training tuple.
+      num_queries: Integer, number of query images to be processed in one epoch.
+      pool_size: Integer, size of the negative image pool, from where the
+        hard-negative images are re-mined.
+      loader: Callable, a function to load an image given its path.
+      ims_root: String, image root directory.
+
+    Raises:
+      ValueError: If mode is not either 'train' or 'val'.
+    """
+
+    if mode not in ['train', 'val']:
+      raise ValueError(
+              "`mode` argument should be either 'train' or 'val', passed as a "
+              "String.")
+
+    # Loading db.
+    db_filename = os.path.join(data_root, '{}.pkl'.format(name))
+    with tf.io.gfile.GFile(db_filename, 'rb') as f:
+      db = pickle.load(f)[mode]
+
+    # Initializing tuples dataset.
+    self._ims_root = data_root if ims_root is None else ims_root
+    self._name = name
+    self._mode = mode
+    self._imsize = imsize
+    self._clusters = db['cluster']
+    self._query_pool = db['qidxs']
+    self._positive_pool = db['pidxs']
+
+    if not hasattr(self, 'images'):
+      self.images = db['ids']
+
+    # Size of training subset for an epoch.
+    self._num_negatives = num_negatives
+    self._num_queries = min(num_queries, len(self._query_pool))
+    self._pool_size = min(pool_size, len(self.images))
+    self._qidxs = None
+    self._pidxs = None
+    self._nidxs = None
+
+    self._loader = loader
+    self._print_freq = 10
+    # Indexer for the iterator.
+    self._n = 0
+
+  def __iter__(self):
+    """Function for making TupleDataset an iterator.
+
+    Returns:
+      iter: The iterator object itself (TupleDataset).
+    """
+    return self
+
+  def __next__(self):
+    """Function for making TupleDataset an iterator.
+
+    Returns:
+      next: The next item in the sequence (next dataset image tuple).
+    """
+    if self._n < len(self._qidxs):
+      result = self.__getitem__(self._n)
+      self._n += 1
+      return result
+    else:
+      raise StopIteration
+
+  def _img_names_to_full_path(self, image_list):
+    """Converts list of image names to the list of full paths to the images.
+
+    Args:
+      image_list: Image names, either a list or a single image path.
+
+    Returns:
+      image_full_paths: List of full paths to the images.
+    """
+    if not isinstance(image_list, list):
+      return os.path.join(self._ims_root, image_list)
+    return [os.path.join(self._ims_root, img_name) for img_name in image_list]
+
+  def __getitem__(self, index):
+    """Called to load an image tuple at the given `index`.
+
+    Args:
+      index: Integer, index.
+
+    Returns:
+      output: Tuple [q,p,n1,...,nN, target], loaded 'train'/'val' tuple at
+        index of qidxs. `q` is the query image tensor, `p` is the
+        corresponding positive image tensor, `n1`,...,`nN` are the negatives
+        associated with the query. `target` is a tensor (with the shape [2+N])
+        of integer labels corresponding to the tuple list: query (-1),
+        positive (1), negative (0).
+
+    Raises:
+      ValueError: Raised if the query indexes list `qidxs` is empty.
+    """
+    if self.__len__() == 0:
+      raise ValueError(
+              "List `qidxs` is empty. Run `dataset.create_epoch_tuples(net)` "
+              "method to create subset for `train`/`val`.")
+
+    output = []
+    # Query image.
+    output.append(self._loader(
+            self._img_names_to_full_path(self.images[self._qidxs[index]]),
+            self._imsize))
+    # Positive image.
+    output.append(self._loader(
+            self._img_names_to_full_path(self.images[self._pidxs[index]]),
+            self._imsize))
+    # Negative images.
+    for nidx in self._nidxs[index]:
+      output.append(self._loader(
+              self._img_names_to_full_path(self.images[nidx]),
+              self._imsize))
+    # Labels for the query (-1), positive (1), negative (0) images in the tuple.
+    target = tf.convert_to_tensor([-1, 1] + [0] * self._num_negatives)
+    output.append(target)
+
+    return tuple(output)
+
+  def __len__(self):
+    """Called to implement the built-in function len().
+
+    Returns:
+      len: Integer, number of query images.
+    """
+    if self._qidxs is None:
+      return 0
+    return len(self._qidxs)
+
+  def __repr__(self):
+    """Metadata for the TupleDataset.
+
+    Returns:
+      meta: String, containing TupleDataset meta.
+    """
+    fmt_str = self.__class__.__name__ + '\n'
+    fmt_str += '\tName and mode: {} {}\n'.format(self._name, self._mode)
+    fmt_str += '\tNumber of images: {}\n'.format(len(self.images))
+    fmt_str += '\tNumber of training tuples: {}\n'.format(len(self._query_pool))
+    fmt_str += '\tNumber of negatives per tuple: {}\n'.format(
+            self._num_negatives)
+    fmt_str += '\tNumber of tuples processed in an epoch: {}\n'.format(
+            self._num_queries)
+    fmt_str += '\tPool size for negative remining: {}\n'.format(self._pool_size)
+    return fmt_str
+
+  def create_epoch_tuples(self, net):
+    """Creates epoch tuples with the hard-negative re-mining.
+
+    Negative examples are selected from clusters different than the cluster
+    of the query image, as the clusters are ideally non-overlaping. For
+    every query image we choose  hard-negatives, that is, non-matching images
+    with the most similar descriptor. Hard-negatives depend on the current
+    CNN parameters. K-nearest neighbors from all non-matching images are
+    selected. Query images are selected randomly. Positives examples are
+    fixed for the related query image during the whole training process.
+
+    Args:
+      net: Model, network to be used for negative re-mining.
+
+    Raises:
+      ValueError: If the pool_size is smaller than the number of negative
+        images per tuple.
+
+    Returns:
+      avg_l2: Float, average negative L2-distance.
+    """
+    self._n = 0
+
+    if self._num_negatives < self._pool_size:
+      raise ValueError("Unable to create epoch tuples. Negative pool_size "
+                       "should be larger than the number of negative images "
+                       "per tuple.")
+
+    global_features_utils.debug_and_log(
+            '>> Creating tuples for an epoch of {}-{}...'.format(self._name,
+                                                                 self._mode),
+            True)
+    global_features_utils.debug_and_log(">> Used network: ", True)
+    global_features_utils.debug_and_log(net.meta_repr(), True)
+
+    ## Selecting queries.
+    # Draw `num_queries` random queries for the tuples.
+    idx_list = np.arange(len(self._query_pool))
+    np.random.shuffle(idx_list)
+    idxs2query_pool = idx_list[:self._num_queries]
+    self._qidxs = [self._query_pool[i] for i in idxs2query_pool]
+
+    ## Selecting positive pairs.
+    # Positives examples are fixed for each query during the whole training
+    # process.
+    self._pidxs = [self._positive_pool[i] for i in idxs2query_pool]
+
+    ## Selecting negative pairs.
+    # If `num_negatives` = 0 create dummy nidxs.
+    # Useful when only positives used for training.
+    if self._num_negatives == 0:
+      self._nidxs = [[] for _ in range(len(self._qidxs))]
+      return 0
+
+    # Draw pool_size random images for pool of negatives images.
+    neg_idx_list = np.arange(len(self.images))
+    np.random.shuffle(neg_idx_list)
+    neg_images_idxs = neg_idx_list[:self._pool_size]
+
+    global_features_utils.debug_and_log(
+            '>> Extracting descriptors for query images...', debug=True)
+
+    img_list = self._img_names_to_full_path([self.images[i] for i in
+                                             self._qidxs])
+    qvecs = global_model.extract_global_descriptors_from_list(
+            net,
+            images=img_list,
+            image_size=self._imsize,
+            print_freq=self._print_freq)
+
+    global_features_utils.debug_and_log(
+            '>> Extracting descriptors for negative pool...', debug=True)
+
+    poolvecs = global_model.extract_global_descriptors_from_list(
+            net,
+            images=self._img_names_to_full_path([self.images[i] for i in
+                                                 neg_images_idxs]),
+            image_size=self._imsize,
+            print_freq=self._print_freq)
+
+    global_features_utils.debug_and_log('>> Searching for hard negatives...',
+                                        debug=True)
+
+    # Compute dot product scores and ranks.
+    scores = tf.linalg.matmul(poolvecs, qvecs, transpose_a=True)
+    ranks = tf.argsort(scores, axis=0, direction='DESCENDING')
+
+    sum_ndist = 0.
+    n_ndist = 0.
+
+    # Selection of negative examples.
+    self._nidxs = []
+
+    for q, qidx in enumerate(self._qidxs):
+      # We are not using the query cluster, those images are potentially
+      # positive.
+      qcluster = self._clusters[qidx]
+      clusters = [qcluster]
+      nidxs = []
+      rank = 0
+
+      while len(nidxs) < self._num_negatives:
+        if rank >= tf.shape(ranks)[0]:
+          raise ValueError("Unable to create epoch tuples. Number of required "
+                           "negative images is larger than the number of "
+                           "clusters in the dataset.")
+        potential = neg_images_idxs[ranks[rank, q]]
+        # Take at most one image from the same cluster.
+        if not self._clusters[potential] in clusters:
+          nidxs.append(potential)
+          clusters.append(self._clusters[potential])
+          dist = tf.norm(qvecs[:, q] - poolvecs[:, ranks[rank, q]],
+                         axis=0).numpy()
+          sum_ndist += dist
+          n_ndist += 1
+        rank += 1
+
+      self._nidxs.append(nidxs)
+
+    global_features_utils.debug_and_log(
+            '>> Average negative l2-distance: {:.2f}'.format(
+                    sum_ndist / n_ndist))
+
+    # Return average negative L2-distance.
+    return sum_ndist / n_ndist

research/delf/delf/python/datasets/tuples_dataset_test.py

+# Lint as: python3
+# Copyright 2021 The TensorFlow 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.
+# ==============================================================================
+"Tests for the tuples dataset module."
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from absl import flags
+import numpy as np
+from PIL import Image
+import tensorflow as tf
+import pickle
+
+from delf.python.datasets import tuples_dataset
+from delf.python.training.model import global_model
+
+FLAGS = flags.FLAGS
+
+
+class TuplesDatasetTest(tf.test.TestCase):
+  """Tests for tuples dataset module."""
+
+  def testCreateEpochTuples(self):
+    """Tests epoch tuple creation."""
+    # Create a tuples dataset instance.
+    name = 'test_dataset'
+    num_queries = 1
+    pool_size = 5
+    num_negatives = 2
+    # Create a ground truth .pkl file.
+    gnd = {
+      'train': {'ids': [str(i) + '.png' for i in range(2 * num_queries + pool_size)],
+                'cluster': [0, 0, 1, 2, 3, 4, 5],
+                'qidxs': [0], 'pidxs': [1]}}
+    gnd_name = name + '.pkl'
+    with tf.io.gfile.GFile(os.path.join(FLAGS.test_tmpdir, gnd_name),
+                           'wb') as gnd_file:
+      pickle.dump(gnd, gnd_file)
+
+    # Create random images for the dataset.
+    for i in range(2 * num_queries + pool_size):
+      dummy_image = np.random.rand(1024, 750, 3) * 255
+      img_out = Image.fromarray(dummy_image.astype('uint8')).convert('RGB')
+      filename = os.path.join(FLAGS.test_tmpdir, '{}.png'.format(i))
+      img_out.save(filename)
+
+    dataset = tuples_dataset.TuplesDataset(
+      name=name,
+      data_root=FLAGS.test_tmpdir,
+      mode='train',
+      imsize=1024,
+      num_negatives=num_negatives,
+      num_queries=num_queries,
+      pool_size=pool_size
+    )
+
+    # Assert that initially no negative images are set.
+    self.assertIsNone(dataset._nidxs)
+
+    # Initialize a network for negative re-mining.
+    model_params = {'architecture': 'ResNet101', 'pooling': 'gem',
+                    'whitening': False, 'pretrained': True}
+    model = global_model.GlobalFeatureNet(**model_params)
+
+    avg_neg_distance = dataset.create_epoch_tuples(model)
+
+    # Check that an appropriate number of negative images has been chosen per
+    # query.
+    self.assertAllEqual(tf.shape(dataset._nidxs), [num_queries, num_negatives])
+
+if __name__ == '__main__':
+  tf.test.main()

research/delf/delf/python/training/global_features_utils.py

-# Copyright 2021 The TensorFlow 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.
-# ==============================================================================
-"""Utilities for the global model training."""
-
-import os
-
-from absl import logging
-
-import numpy as np
-import tensorflow as tf
-
-from delf.python.datasets.revisited_op import dataset
-
-
-class AverageMeter():
-  """Computes and stores the average and current value of loss."""
-
-  def __init__(self):
-    """Initialization of the AverageMeter."""
-    self.reset()
-
-  def reset(self):
-    """Resets all the values."""
-    self.val = 0
-    self.avg = 0
-    self.sum = 0
-    self.count = 0
-
-  def update(self, val, n=1):
-    """Updates values in the AverageMeter.
-
-    Args:
-      val: Float, loss value.
-      n: Integer, number of instances.
-    """
-    self.val = val
-    self.sum += val * n
-    self.count += n
-    self.avg = self.sum / self.count
-
-
-def compute_metrics_and_print(dataset_name,
-                              sorted_index_ids,
-                              ground_truth,
-                              desired_pr_ranks=None,
-                              log=True):
-  """Computes and logs ground-truth metrics for Revisited datasets.
-
-  Args:
-    dataset_name: String, name of the dataset.
-    sorted_index_ids: Integer NumPy array of shape [#queries, #index_images].
-      For each query, contains an array denoting the most relevant index images,
-      sorted from most to least relevant.
-    ground_truth: List containing ground-truth information for dataset. Each
-      entry is a dict corresponding to the ground-truth information for a query.
-      The dict has keys 'ok' and 'junk', mapping to a NumPy array of integers.
-    desired_pr_ranks: List of integers containing the desired precision/recall
-      ranks to be reported. E.g., if precision@1/recall@1 and
-      precision@10/recall@10 are desired, this should be set to [1, 10]. The
-      largest item should be <= #sorted_index_ids. Default: [1, 5, 10].
-    log: Whether to log results using logging.info().
-
-  Returns:
-    mAP: (metricsE, metricsM, metricsH) Tuple of the metrics for different
-      levels of complexity. Each metrics is a list containing:
-      mean_average_precision (float), mean_precisions (NumPy array of
-      floats, with shape [len(desired_pr_ranks)]), mean_recalls (NumPy array
-      of floats, with shape [len(desired_pr_ranks)]), average_precisions
-      (NumPy array of floats, with shape [#queries]), precisions (NumPy array of
-      floats, with shape [#queries, len(desired_pr_ranks)]), recalls (NumPy
-      array of floats, with shape [#queries, len(desired_pr_ranks)]).
-
-  Raises:
-    ValueError: If an unknown dataset name is provided as an argument.
-  """
-  if dataset not in dataset.DATASET_NAMES:
-    raise ValueError('Unknown dataset: {}!'.format(dataset))
-
-  if desired_pr_ranks is None:
-    desired_pr_ranks = [1, 5, 10]
-
-  (easy_ground_truth, medium_ground_truth,
-   hard_ground_truth) = dataset.ParseEasyMediumHardGroundTruth(ground_truth)
-
-  metrics_easy = dataset.ComputeMetrics(sorted_index_ids, easy_ground_truth,
-                                        desired_pr_ranks)
-  metrics_medium = dataset.ComputeMetrics(sorted_index_ids, medium_ground_truth,
-                                          desired_pr_ranks)
-  metrics_hard = dataset.ComputeMetrics(sorted_index_ids, hard_ground_truth,
-                                        desired_pr_ranks)
-
-  debug_and_log(
-      '>> {}: mAP E: {}, M: {}, H: {}'.format(
-          dataset_name, np.around(metrics_easy[0] * 100, decimals=2),
-          np.around(metrics_medium[0] * 100, decimals=2),
-          np.around(metrics_hard[0] * 100, decimals=2)),
-      log=log)
-
-  debug_and_log(
-      '>> {}: mP@k{} E: {}, M: {}, H: {}'.format(
-          dataset_name, desired_pr_ranks,
-          np.around(metrics_easy[1] * 100, decimals=2),
-          np.around(metrics_medium[1] * 100, decimals=2),
-          np.around(metrics_hard[1] * 100, decimals=2)),
-      log=log)
-
-  return metrics_easy, metrics_medium, metrics_hard
-
-
-def htime(time_difference):
-  """Time formatting function.
-
-  Depending on the value of `time_difference` outputs time in an appropriate
-  time format.
-
-  Args:
-    time_difference: Float, time difference between the two events.
-
-  Returns:
-    time: String representing time in an appropriate time format.
-  """
-  time_difference = round(time_difference)
-
-  days = time_difference // 86400
-  hours = time_difference // 3600 % 24
-  minutes = time_difference // 60 % 60
-  seconds = time_difference % 60
-
-  if days > 0:
-    return '{:d}d {:d}h {:d}m {:d}s'.format(days, hours, minutes, seconds)
-  if hours > 0:
-    return '{:d}h {:d}m {:d}s'.format(hours, minutes, seconds)
-  if minutes > 0:
-    return '{:d}m {:d}s'.format(minutes, seconds)
-  return '{:d}s'.format(seconds)
-
-
-def debug_and_log(msg, debug=True, log=True, debug_on_the_same_line=False):
-  """Outputs `msg` to both stdout (if in the debug mode) and the log file.
-
-  Args:
-    msg: String, message to be logged.
-    debug: Bool, if True, will print `msg` to stdout.
-    log: Bool, if True, will redirect `msg` to the logfile.
-    debug_on_the_same_line: Bool, if True, will print `msg` to stdout without a
-      new line. When using this mode, logging to a logfile is disabled.
-  """
-  if debug_on_the_same_line:
-    print(msg, end='')
-    return
-  if debug:
-    print(msg)
-  if log:
-    logging.info(msg)
-
-
-def get_standard_keras_models():
-  """Gets the standard keras model names.
-
-  Returns:
-    model_names: List, names of the standard keras models.
-  """
-  model_names = sorted(
-      name for name in tf.keras.applications.__dict__
-      if not name.startswith('__') and
-      callable(tf.keras.applications.__dict__[name]))
-  return model_names
-
-
-def create_model_directory(training_dataset, arch, pool, whitening, pretrained,
-                           loss, loss_margin, optimizer, lr, weight_decay,
-                           neg_num, query_size, pool_size, batch_size,
-                           update_every, image_size, directory):
-  """Based on the model parameters, creates the model directory.
-
-  If the model directory does not exist, the directory is created.
-
-  Args:
-    training_dataset: String, training dataset name.
-    arch: String, model architecture.
-    pool: String, pooling option.
-    whitening: Bool, whether the model is trained with global whitening.
-    pretrained: Bool, whether the model is initialized with the precomputed
-      weights.
-    loss: String, training loss type.
-    loss_margin: Float, loss margin.
-    optimizer: Sting, used optimizer.
-    lr: Float, initial learning rate.
-    weight_decay: Float, weight decay.
-    neg_num: Integer, Number of negative images per train/val tuple.
-    query_size: Integer, number of queries per one training epoch.
-    pool_size: Integer, size of the pool for hard negative mining.
-    batch_size: Integer, batch size.
-    update_every: Integer, frequency of the model weights update.
-    image_size: Integer, maximum size of longer image side used for training.
-    directory: String, destination where trained network should be saved.
-
-  Returns:
-    folder: String, path to the model folder.
-  """
-  folder = '{}_{}_{}'.format(training_dataset, arch, pool)
-  if whitening:
-    folder += '_whiten'
-  if not pretrained:
-    folder += '_notpretrained'
-  folder += ('_{}_m{:.2f}_{}_lr{:.1e}_wd{:.1e}_nnum{}_qsize{}_psize{}_bsize{}'
-             '_uevery{}_imsize{}').format(loss, loss_margin, optimizer, lr,
-                                          weight_decay, neg_num, query_size,
-                                          pool_size, batch_size, update_every,
-                                          image_size)
-
-  folder = os.path.join(directory, folder)
-  debug_and_log(
-      '>> Creating directory if does not exist:\n>> \'{}\''.format(folder))
-  if not os.path.exists(folder):
-    os.makedirs(folder)
-  return folder