


"""simple docstring"""







import math
from typing import Dict, Iterable, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    is_torch_available,
    is_torch_tensor,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_torch_available():
     import torch

if is_vision_available():
     import PIL


lowerCAmelCase__      				=   logging.get_logger(__name__)
def 					a__			(				SCREAMING_SNAKE_CASE    :    np.ndarray				,      SCREAMING_SNAKE_CASE    :    Union[int, Iterable[int]]				,      SCREAMING_SNAKE_CASE    :    bool				,      SCREAMING_SNAKE_CASE    :    int				):


  '''simple docstring'''







  def constraint_to_multiple_of(SCREAMING_SNAKE_CASE    :    List[Any]				,      SCREAMING_SNAKE_CASE    :    Any				,      SCREAMING_SNAKE_CASE    :    Union[str, Any]=0				,      SCREAMING_SNAKE_CASE    :    Optional[int]=None				):
    lowerCAmelCase     :     List[Any]								=							round(val / multiple				) * multiple

    if max_val is not None and x > max_val:
      lowerCAmelCase     :     int								=							math.floor(val / multiple				) * multiple

    if x < min_val:
      lowerCAmelCase     :     Optional[int]								=							math.ceil(val / multiple				) * multiple

    return x

  lowerCAmelCase     :     List[Any]								=							(output_size, output_size) if isinstance(SCREAMING_SNAKE_CASE				,      SCREAMING_SNAKE_CASE				) else output_size

  lowerCAmelCase						,  lowerCAmelCase     :     Tuple								=							get_image_size(SCREAMING_SNAKE_CASE				)
  lowerCAmelCase						,  lowerCAmelCase     :     Optional[Any]								=							output_size

  # determine new height and width
  lowerCAmelCase     :     Dict								=							output_height / input_height
  lowerCAmelCase     :     Tuple								=							output_width / input_width

  if keep_aspect_ratio:
    # scale as little as possible
    if abs(1 - scale_width				) < abs(1 - scale_height				):
      # fit width
      lowerCAmelCase     :     Dict								=							scale_width
    else:
      # fit height
      lowerCAmelCase     :     Dict								=							scale_height

  lowerCAmelCase     :     List[Any]								=							constraint_to_multiple_of(scale_height * input_height				,      multiple=SCREAMING_SNAKE_CASE				)
  lowerCAmelCase     :     Optional[Any]								=							constraint_to_multiple_of(scale_width * input_width				,      multiple=SCREAMING_SNAKE_CASE				)

  return (new_height, new_width)





class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       a  :				Any   =["pixel_values"]

       def __init__(						self		,     snake_case__ = True		,     snake_case__ = None		,     snake_case__ = PILImageResampling.BILINEAR		,     snake_case__ = False		,     snake_case__ = 1		,     snake_case__ = True		,     snake_case__ = 1 / 255		,     snake_case__ = True		,     snake_case__ = None		,     snake_case__ = None		,     **snake_case__		,     ):






         """simple docstring"""





         super().__init__(**snake_case__						)
         lowerCAmelCase     :     int								=							size if size is not None else {"height": 384, "width": 384}
         lowerCAmelCase     :     Tuple								=							get_size_dict(snake_case__						)
         lowerCAmelCase     :     int								=							do_resize
         lowerCAmelCase     :     Tuple								=							size
         lowerCAmelCase     :     Dict								=							keep_aspect_ratio
         lowerCAmelCase     :     List[Any]								=							ensure_multiple_of
         lowerCAmelCase     :     Optional[int]								=							resample
         lowerCAmelCase     :     Tuple								=							do_rescale
         lowerCAmelCase     :     int								=							rescale_factor
         lowerCAmelCase     :     List[Any]								=							do_normalize
         lowerCAmelCase     :     str								=							image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
         lowerCAmelCase     :     List[str]								=							image_std if image_std is not None else IMAGENET_STANDARD_STD

       def        lowercase__							(						self		,     snake_case__		,     snake_case__		,     snake_case__ = False		,     snake_case__ = 1		,     snake_case__ = PILImageResampling.BICUBIC		,     snake_case__ = None		,     **snake_case__		,     ):






         """simple docstring"""





         lowerCAmelCase     :     int								=							get_size_dict(snake_case__						)
         if "height" not in size or "width" not in size:
           raise ValueError(f"""The size dictionary must contain the keys 'height' and 'width'. Got {size.keys()}"""						)
         lowerCAmelCase     :     int								=							get_resize_output_image_size(
             snake_case__		,     output_size=(size["height"], size["width"])		,     keep_aspect_ratio=snake_case__		,     multiple=snake_case__		,     )
         return resize(snake_case__		,     size=snake_case__		,     resample=snake_case__		,     data_format=snake_case__		,     **snake_case__						)

       def        lowercase__							(						self		,     snake_case__		,     snake_case__		,     snake_case__ = None		,     **snake_case__		,     ):






         """simple docstring"""





         return rescale(snake_case__		,     scale=snake_case__		,     data_format=snake_case__		,     **snake_case__						)

       def        lowercase__							(						self		,     snake_case__		,     snake_case__		,     snake_case__		,     snake_case__ = None		,     **snake_case__		,     ):






         """simple docstring"""





         return normalize(snake_case__		,     mean=snake_case__		,     std=snake_case__		,     data_format=snake_case__		,     **snake_case__						)

       def        lowercase__							(						self		,     snake_case__		,     snake_case__ = None		,     snake_case__ = None		,     snake_case__ = None		,     snake_case__ = None		,     snake_case__ = None		,     snake_case__ = None		,     snake_case__ = None		,     snake_case__ = None		,     snake_case__ = None		,     snake_case__ = None		,     snake_case__ = None		,     snake_case__ = ChannelDimension.FIRST		,     **snake_case__		,     ):






         """simple docstring"""





         lowerCAmelCase     :     Tuple								=							do_resize if do_resize is not None else self.do_resize
         lowerCAmelCase     :     Dict								=							size if size is not None else self.size
         lowerCAmelCase     :     Any								=							get_size_dict(snake_case__						)
         lowerCAmelCase     :     Any								=							keep_aspect_ratio if keep_aspect_ratio is not None else self.keep_aspect_ratio
         lowerCAmelCase     :     List[Any]								=							ensure_multiple_of if ensure_multiple_of is not None else self.ensure_multiple_of
         lowerCAmelCase     :     Tuple								=							resample if resample is not None else self.resample
         lowerCAmelCase     :     List[Any]								=							do_rescale if do_rescale is not None else self.do_rescale
         lowerCAmelCase     :     str								=							rescale_factor if rescale_factor is not None else self.rescale_factor
         lowerCAmelCase     :     Tuple								=							do_normalize if do_normalize is not None else self.do_normalize
         lowerCAmelCase     :     str								=							image_mean if image_mean is not None else self.image_mean
         lowerCAmelCase     :     Tuple								=							image_std if image_std is not None else self.image_std

         lowerCAmelCase     :     str								=							make_list_of_images(snake_case__						)

         if not valid_images(snake_case__						):
           raise ValueError(
               "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
               "torch.Tensor, tf.Tensor or jax.ndarray."						)

         if do_resize and size is None or resample is None:
           raise ValueError("Size and resample must be specified if do_resize is True."						)

         if do_rescale and rescale_factor is None:
           raise ValueError("Rescale factor must be specified if do_rescale is True."						)

         if do_normalize and (image_mean is None or image_std is None):
           raise ValueError("Image mean and std must be specified if do_normalize is True."						)

         # All transformations expect numpy arrays.
         lowerCAmelCase     :     Optional[Any]								=							[to_numpy_array(snake_case__						) for image in images]

         if do_resize:
           lowerCAmelCase     :     Dict								=							[self.resize(image=snake_case__		,     size=snake_case__		,     resample=snake_case__						) for image in images]

         if do_rescale:
           lowerCAmelCase     :     Optional[int]								=							[self.rescale(image=snake_case__		,     scale=snake_case__						) for image in images]

         if do_normalize:
           lowerCAmelCase     :     Optional[int]								=							[self.normalize(image=snake_case__		,     mean=snake_case__		,     std=snake_case__						) for image in images]

         lowerCAmelCase     :     Optional[Any]								=							[to_channel_dimension_format(snake_case__		,     snake_case__						) for image in images]

         lowerCAmelCase     :     Any								=							{"pixel_values": images}
         return BatchFeature(data=snake_case__		,     tensor_type=snake_case__						)







       def        lowercase__							(						self		,     snake_case__		,     snake_case__ = None						):






         """simple docstring"""





         lowerCAmelCase     :     Optional[int]								=							outputs.logits

         # Resize logits and compute semantic segmentation maps
         if target_sizes is not None:
           if len(snake_case__						) != len(snake_case__						):
             raise ValueError(
                 "Make sure that you pass in as many target sizes as the batch dimension of the logits"						)

           if is_torch_tensor(snake_case__						):
             lowerCAmelCase     :     Optional[Any]								=							target_sizes.numpy()

           lowerCAmelCase     :     Dict								=							[]

           for idx in range(len(snake_case__						)						):
             lowerCAmelCase     :     int								=							torch.nn.functional.interpolate(
                 logits[idx].unsqueeze(dim=0						)		,     size=target_sizes[idx]		,     mode="bilinear"		,     align_corners=snake_case__						)
             lowerCAmelCase     :     List[str]								=							resized_logits[0].argmax(dim=0						)
             semantic_segmentation.append(snake_case__						)
         else:
           lowerCAmelCase     :     int								=							logits.argmax(dim=1						)
           lowerCAmelCase     :     Optional[Any]								=							[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]						)]

         return semantic_segmentation









"""simple docstring"""







import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pytest

from datasets.arrow_dataset import Dataset
from datasets.search import ElasticSearchIndex, FaissIndex, MissingIndex

from .utils import require_elasticsearch, require_faiss


lowerCAmelCase__      				=   pytest.mark.integration


@require_faiss
class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     List[str]								=							Dataset.from_dict({"filename": ["my_name-train" + "_" + str(snake_case__						) for x in np.arange(30						).tolist()]}						)
         return dset

       def        lowercase__							(						self						):






         """simple docstring"""





         import faiss

         lowerCAmelCase     :     Dataset								=							self._create_dummy_dataset()
         lowerCAmelCase     :     Union[str, Any]								=							dset.map(
             lambda snake_case__		,     snake_case__						: {"vecs": i * np.ones(5		,     dtype=np.floataa						)}		,     with_indices=snake_case__		,     keep_in_memory=snake_case__						)
         lowerCAmelCase     :     Union[str, Any]								=							dset.add_faiss_index("vecs"		,     batch_size=100		,     metric_type=faiss.METRIC_INNER_PRODUCT						)
         lowerCAmelCase						,  lowerCAmelCase     :     Tuple								=							dset.get_nearest_examples("vecs"		,     np.ones(5		,     dtype=np.floataa						)						)
         self.assertEqual(examples["filename"][0]		,     "my_name-train_29"						)
         dset.drop_index("vecs"						)

       def        lowercase__							(						self						):






         """simple docstring"""





         import faiss

         lowerCAmelCase     :     Dataset								=							self._create_dummy_dataset()
         dset.add_faiss_index_from_external_arrays(
             external_arrays=np.ones((30, 5)						) * np.arange(30						).reshape(-1		,     1						)		,     index_name="vecs"		,     batch_size=100		,     metric_type=faiss.METRIC_INNER_PRODUCT		,     )
         lowerCAmelCase						,  lowerCAmelCase     :     Optional[Any]								=							dset.get_nearest_examples("vecs"		,     np.ones(5		,     dtype=np.floataa						)						)
         self.assertEqual(examples["filename"][0]		,     "my_name-train_29"						)

       def        lowercase__							(						self						):






         """simple docstring"""





         import faiss

         lowerCAmelCase     :     Dataset								=							self._create_dummy_dataset()
         dset.add_faiss_index_from_external_arrays(
             external_arrays=np.ones((30, 5)						) * np.arange(30						).reshape(-1		,     1						)		,     index_name="vecs"		,     metric_type=faiss.METRIC_INNER_PRODUCT		,     )

         # Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
         # ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
         # see https://bugs.python.org/issue14243 and
         # https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
         with tempfile.NamedTemporaryFile(delete=snake_case__						) as tmp_file:
           dset.save_faiss_index("vecs"		,     tmp_file.name						)
           dset.load_faiss_index("vecs2"		,     tmp_file.name						)
         os.unlink(tmp_file.name						)

         lowerCAmelCase						,  lowerCAmelCase     :     int								=							dset.get_nearest_examples("vecs2"		,     np.ones(5		,     dtype=np.floataa						)						)
         self.assertEqual(examples["filename"][0]		,     "my_name-train_29"						)

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     Dataset								=							self._create_dummy_dataset()
         dset.add_faiss_index_from_external_arrays(
             external_arrays=np.ones((30, 5)						) * np.arange(30						).reshape(-1		,     1						)		,     index_name="vecs"						)
         dset.drop_index("vecs"						)
         self.assertRaises(snake_case__		,     partial(dset.get_nearest_examples		,     "vecs2"		,     np.ones(5		,     dtype=np.floataa						)						)						)







       def        lowercase__							(						self						):






         """simple docstring"""





         from elasticsearch import Elasticsearch

         lowerCAmelCase     :     Dataset								=							self._create_dummy_dataset()
         with patch("elasticsearch.Elasticsearch.search"						) as mocked_search, patch(
             "elasticsearch.client.IndicesClient.create"						) as mocked_index_create, patch("elasticsearch.helpers.streaming_bulk"						) as mocked_bulk:
           lowerCAmelCase     :     List[str]								=							{"acknowledged": True}
           mocked_bulk.return_value([(True, None)] * 30						)
           lowerCAmelCase     :     List[str]								=							{"hits": {"hits": [{"_score": 1, "_id": 29}]}}
           lowerCAmelCase     :     str								=							Elasticsearch()

           dset.add_elasticsearch_index("filename"		,     es_client=snake_case__						)
           lowerCAmelCase						,  lowerCAmelCase     :     int								=							dset.get_nearest_examples("filename"		,     "my_name-train_29"						)
           self.assertEqual(examples["filename"][0]		,     "my_name-train_29"						)


@require_faiss
class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       def        lowercase__							(						self						):






         """simple docstring"""





         import faiss

         lowerCAmelCase     :     Optional[int]								=							FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT						)

         # add vectors
         index.add_vectors(np.eye(5		,     dtype=np.floataa						)						)
         self.assertIsNotNone(index.faiss_index						)
         self.assertEqual(index.faiss_index.ntotal		,     5						)
         index.add_vectors(np.zeros((5, 5)		,     dtype=np.floataa						)						)
         self.assertEqual(index.faiss_index.ntotal		,     10						)

         # single query
         lowerCAmelCase     :     int								=							np.zeros(5		,     dtype=np.floataa						)
         lowerCAmelCase     :     Optional[int]								=							1
         lowerCAmelCase						,  lowerCAmelCase     :     Optional[Any]								=							index.search(snake_case__						)
         self.assertRaises(snake_case__		,     index.search		,     query.reshape(-1		,     1						)						)
         self.assertGreater(scores[0]		,     0						)
         self.assertEqual(indices[0]		,     1						)

         # batched queries
         lowerCAmelCase     :     Union[str, Any]								=							np.eye(5		,     dtype=np.floataa						)[::-1]
         lowerCAmelCase						,  lowerCAmelCase     :     str								=							index.search_batch(snake_case__						)
         self.assertRaises(snake_case__		,     index.search_batch		,     queries[0]						)
         lowerCAmelCase     :     Optional[int]								=							[scores[0] for scores in total_scores]
         lowerCAmelCase     :     int								=							[indices[0] for indices in total_indices]
         self.assertGreater(np.min(snake_case__						)		,     0						)
         self.assertListEqual([4, 3, 2, 1, 0]		,     snake_case__						)

       def        lowercase__							(						self						):






         """simple docstring"""





         import faiss

         lowerCAmelCase     :     Dict								=							FaissIndex(string_factory="Flat"						)
         index.add_vectors(np.eye(5		,     dtype=np.floataa						)						)
         self.assertIsInstance(index.faiss_index		,     faiss.IndexFlat						)
         lowerCAmelCase     :     Union[str, Any]								=							FaissIndex(string_factory="LSH"						)
         index.add_vectors(np.eye(5		,     dtype=np.floataa						)						)
         self.assertIsInstance(index.faiss_index		,     faiss.IndexLSH						)
         with self.assertRaises(snake_case__						):
           lowerCAmelCase     :     List[Any]								=							FaissIndex(string_factory="Flat"		,     custom_index=faiss.IndexFlat(5						)						)

       def        lowercase__							(						self						):






         """simple docstring"""





         import faiss

         lowerCAmelCase     :     Any								=							faiss.IndexFlat(5						)
         lowerCAmelCase     :     Union[str, Any]								=							FaissIndex(custom_index=snake_case__						)
         index.add_vectors(np.eye(5		,     dtype=np.floataa						)						)
         self.assertIsInstance(index.faiss_index		,     faiss.IndexFlat						)







       def        lowercase__							(						self						):






         """simple docstring"""





         import faiss

         lowerCAmelCase     :     str								=							FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT						)
         index.add_vectors(np.eye(5		,     dtype=np.floataa						)						)

         # Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
         # ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
         # see https://bugs.python.org/issue14243 and
         # https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
         with tempfile.NamedTemporaryFile(delete=snake_case__						) as tmp_file:
           index.save(tmp_file.name						)
           lowerCAmelCase     :     int								=							FaissIndex.load(tmp_file.name						)
         os.unlink(tmp_file.name						)

         lowerCAmelCase     :     Union[str, Any]								=							np.zeros(5		,     dtype=np.floataa						)
         lowerCAmelCase     :     List[str]								=							1
         lowerCAmelCase						,  lowerCAmelCase     :     Tuple								=							index.search(snake_case__						)
         self.assertGreater(scores[0]		,     0						)
         self.assertEqual(indices[0]		,     1						)

@require_faiss
def 					a__			(				SCREAMING_SNAKE_CASE    :    Any				):


  '''simple docstring'''







  import faiss

  lowerCAmelCase     :     List[Any]								=							FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT				)
  index.add_vectors(np.eye(5				,      dtype=np.floataa				)				)

  lowerCAmelCase     :     Union[str, Any]								=							"index.faiss"
  lowerCAmelCase     :     List[str]								=							f"""mock://{index_name}"""
  index.save(SCREAMING_SNAKE_CASE				,      storage_options=mockfs.storage_options				)
  lowerCAmelCase     :     Optional[Any]								=							FaissIndex.load(SCREAMING_SNAKE_CASE				,      storage_options=mockfs.storage_options				)

  lowerCAmelCase     :     Optional[int]								=							np.zeros(5				,      dtype=np.floataa				)
  lowerCAmelCase     :     Any								=							1
  lowerCAmelCase						,  lowerCAmelCase     :     Optional[int]								=							index.search(SCREAMING_SNAKE_CASE				)
  assert scores[0] > 0
  assert indices[0] == 1





@require_elasticsearch
class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       def        lowercase__							(						self						):






         """simple docstring"""





         from elasticsearch import Elasticsearch

         with patch("elasticsearch.Elasticsearch.search"						) as mocked_search, patch(
             "elasticsearch.client.IndicesClient.create"						) as mocked_index_create, patch("elasticsearch.helpers.streaming_bulk"						) as mocked_bulk:
           lowerCAmelCase     :     List[str]								=							Elasticsearch()
           lowerCAmelCase     :     Dict								=							{"acknowledged": True}
           lowerCAmelCase     :     Optional[int]								=							ElasticSearchIndex(es_client=snake_case__						)
           mocked_bulk.return_value([(True, None)] * 3						)
           index.add_documents(["foo", "bar", "foobar"]						)

           # single query
           lowerCAmelCase     :     List[str]								=							"foo"
           lowerCAmelCase     :     List[str]								=							{"hits": {"hits": [{"_score": 1, "_id": 0}]}}
           lowerCAmelCase						,  lowerCAmelCase     :     Optional[int]								=							index.search(snake_case__						)
           self.assertEqual(scores[0]		,     1						)
           self.assertEqual(indices[0]		,     0						)

           # single query with timeout
           lowerCAmelCase     :     int								=							"foo"
           lowerCAmelCase     :     Any								=							{"hits": {"hits": [{"_score": 1, "_id": 0}]}}
           lowerCAmelCase						,  lowerCAmelCase     :     str								=							index.search(snake_case__		,     request_timeout=30						)
           self.assertEqual(scores[0]		,     1						)
           self.assertEqual(indices[0]		,     0						)

           # batched queries
           lowerCAmelCase     :     Any								=							["foo", "bar", "foobar"]
           lowerCAmelCase     :     Optional[int]								=							{"hits": {"hits": [{"_score": 1, "_id": 1}]}}
           lowerCAmelCase						,  lowerCAmelCase     :     Any								=							index.search_batch(snake_case__						)
           lowerCAmelCase     :     Tuple								=							[scores[0] for scores in total_scores]
           lowerCAmelCase     :     Dict								=							[indices[0] for indices in total_indices]
           self.assertGreater(np.min(snake_case__						)		,     0						)
           self.assertListEqual([1, 1, 1]		,     snake_case__						)

           # batched queries with timeout
           lowerCAmelCase     :     Optional[Any]								=							["foo", "bar", "foobar"]
           lowerCAmelCase     :     List[Any]								=							{"hits": {"hits": [{"_score": 1, "_id": 1}]}}
           lowerCAmelCase						,  lowerCAmelCase     :     Any								=							index.search_batch(snake_case__		,     request_timeout=30						)
           lowerCAmelCase     :     Dict								=							[scores[0] for scores in total_scores]
           lowerCAmelCase     :     int								=							[indices[0] for indices in total_indices]
           self.assertGreater(np.min(snake_case__						)		,     0						)
           self.assertListEqual([1, 1, 1]		,     snake_case__						)





