






"""simple docstring"""




from typing import Dict, List, Optional, Union

import numpy as np

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


if is_vision_available():
					import PIL


UpperCamelCase							       =    logging.get_logger(__name__)




class    lowercase_    (_UpperCAmelCase     ):
						A__	:		Dict       			=				['''pixel_values''']






						def __init__(    self     ,     a_ = True     ,     a_ = None     ,     a_ = None     ,     a_ = PILImageResampling.BILINEAR     ,     a_ = True     ,     a_ = 1 / 2_5_5     ,     a_ = True     ,     a_ = None     ,     a_ = None     ,     **a_     ,     )	->None:







								'''simple docstring'''


								super().__init__(**a_	)
								_a							=	size if size is not None else {"shortest_edge": 3_8_4}
								_a							=	get_size_dict(a_     ,     default_to_square=a_	)

								_a							=	do_resize
								_a							=	size
								# Default value set here for backwards compatibility where the value in config is None
								_a							=	crop_pct if crop_pct is not None else 2_2_4 / 2_5_6
								_a							=	resample
								_a							=	do_rescale
								_a							=	rescale_factor
								_a							=	do_normalize
								_a							=	image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
								_a							=	image_std if image_std is not None else IMAGENET_STANDARD_STD






						def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_     ,     a_ = PILImageResampling.BICUBIC     ,     a_ = None     ,     **a_     ,     )	->np.ndarray:







								'''simple docstring'''


								_a							=	get_size_dict(a_     ,     default_to_square=a_	)
								if "shortest_edge" not in size:
										raise ValueError(f'''Size dictionary must contain \'shortest_edge\' key. Got {size.keys()}'''	)
								_a							=	size["shortest_edge"]

								if shortest_edge < 3_8_4:
										# maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
										_a							=	int(shortest_edge / crop_pct	)
										_a							=	get_resize_output_image_size(a_     ,     size=a_     ,     default_to_square=a_	)
										_a							=	resize(image=a_     ,     size=a_     ,     resample=a_     ,     data_format=a_     ,     **a_	)
										# then crop to (shortest_edge, shortest_edge)
										return center_crop(image=a_     ,     size=(shortest_edge, shortest_edge)     ,     data_format=a_     ,     **a_	)
								else:
										# warping (no cropping) when evaluated at 384 or larger
										return resize(
										    a_     ,     size=(shortest_edge, shortest_edge)     ,     resample=a_     ,     data_format=a_     ,     **a_	)






						def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_ = None     ,     **a_     ,     )	->List[Any]:







								'''simple docstring'''


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






						def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_     ,     a_ = None     ,     **a_     ,     )	->np.ndarray:







								'''simple docstring'''


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






						def 			lowerCamelCase__    (    self     ,     a_     ,     a_ = None     ,     a_ = None     ,     a_ = None     ,     a_ = None     ,     a_ = None     ,     a_ = None     ,     a_ = None     ,     a_ = None     ,     a_ = None     ,     a_ = None     ,     a_ = ChannelDimension.FIRST     ,     **a_     ,     )	->PIL.Image.Image:







								'''simple docstring'''


								_a							=	do_resize if do_resize is not None else self.do_resize
								_a							=	crop_pct if crop_pct is not None else self.crop_pct
								_a							=	resample if resample is not None else self.resample
								_a							=	do_rescale if do_rescale is not None else self.do_rescale
								_a							=	rescale_factor if rescale_factor is not None else self.rescale_factor
								_a							=	do_normalize if do_normalize is not None else self.do_normalize
								_a							=	image_mean if image_mean is not None else self.image_mean
								_a							=	image_std if image_std is not None else self.image_std

								_a							=	size if size is not None else self.size
								_a							=	get_size_dict(a_     ,     default_to_square=a_	)

								_a							=	make_list_of_images(a_	)

								if not valid_images(a_	):
										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_resize and size["shortest_edge"] < 3_8_4 and crop_pct is None:
										raise ValueError("crop_pct must be specified if size < 384."	)

								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.
								_a							=	[to_numpy_array(a_	) for image in images]

								if do_resize:
										_a							=	[self.resize(image=a_     ,     size=a_     ,     crop_pct=a_     ,     resample=a_	) for image in images]

								if do_rescale:
										_a							=	[self.rescale(image=a_     ,     scale=a_	) for image in images]

								if do_normalize:
										_a							=	[self.normalize(image=a_     ,     mean=a_     ,     std=a_	) for image in images]

								_a							=	[to_channel_dimension_format(a_     ,     a_	) for image in images]

								_a							=	{"pixel_values": images}
								return BatchFeature(data=a_     ,     tensor_type=a_	)










"""simple docstring"""




import os
from typing import BinaryIO, Optional, Union

import numpy as np
import pyarrow.parquet as pq

from .. import Audio, Dataset, Features, Image, NamedSplit, Value, config
from ..features.features import FeatureType, _visit
from ..formatting import query_table
from ..packaged_modules import _PACKAGED_DATASETS_MODULES
from ..packaged_modules.parquet.parquet import Parquet
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader






def        lowerCAmelCase (      UpperCamelCase_:							Features					)       ->							Optional[int]:







  '''simple docstring'''







  _a							=	np.inf

  def set_batch_size(UpperCamelCase_:							FeatureType					) -> None:
    nonlocal batch_size
    if isinstance(UpperCamelCase_     , UpperCamelCase_					):
      _a							=	min(UpperCamelCase_     , config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS					)
    elif isinstance(UpperCamelCase_     , UpperCamelCase_					):
      _a							=	min(UpperCamelCase_     , config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS					)
    elif isinstance(UpperCamelCase_     , UpperCamelCase_					) and feature.dtype == "binary":
      _a							=	min(UpperCamelCase_     , config.PARQUET_ROW_GROUP_SIZE_FOR_BINARY_DATASETS					)

  _visit(UpperCamelCase_     , UpperCamelCase_					)

  return None if batch_size is np.inf else batch_size




class    lowercase_    (_UpperCAmelCase     ):






      def __init__(    self     ,     a_     ,     a_ = None     ,     a_ = None     ,     a_ = None     ,     a_ = False     ,     a_ = False     ,     a_ = None     ,     **a_     ,     )	->Optional[int]:







        '''simple docstring'''


        super().__init__(
            a_     ,     split=a_     ,     features=a_     ,     cache_dir=a_     ,     keep_in_memory=a_     ,     streaming=a_     ,     num_proc=a_     ,     **a_     ,     )
        _a							=	path_or_paths if isinstance(a_     ,     a_	) else {self.split: path_or_paths}
        _a							=	_PACKAGED_DATASETS_MODULES["parquet"][1]
        _a							=	Parquet(
            cache_dir=a_     ,     data_files=a_     ,     features=a_     ,     hash=a_     ,     **a_     ,     )






      def 			lowerCamelCase__    (    self	)	->str:







        '''simple docstring'''


        if self.streaming:
          _a							=	self.builder.as_streaming_dataset(split=self.split	)
        # Build regular (map-style) dataset
        else:
          _a							=	None
          _a							=	None
          _a							=	None
          _a							=	None

          self.builder.download_and_prepare(
              download_config=a_     ,     download_mode=a_     ,     verification_mode=a_     ,     base_path=a_     ,     num_proc=self.num_proc     ,     )
          _a							=	self.builder.as_dataset(
              split=self.split     ,     verification_mode=a_     ,     in_memory=self.keep_in_memory	)
        return dataset







class    lowercase_    :






      def __init__(    self     ,     a_     ,     a_     ,     a_ = None     ,     **a_     ,     )	->int:







        '''simple docstring'''


        _a							=	dataset
        _a							=	path_or_buf
        _a							=	batch_size or get_writer_batch_size(dataset.features	)
        _a							=	parquet_writer_kwargs






      def 			lowerCamelCase__    (    self	)	->int:







        '''simple docstring'''


        _a							=	self.batch_size if self.batch_size else config.DEFAULT_MAX_BATCH_SIZE

        if isinstance(self.path_or_buf     ,     (str, bytes, os.PathLike)	):
          with open(self.path_or_buf     ,     "wb+"	) as buffer:
            _a							=	self._write(file_obj=a_     ,     batch_size=a_     ,     **self.parquet_writer_kwargs	)
        else:
          _a							=	self._write(file_obj=self.path_or_buf     ,     batch_size=a_     ,     **self.parquet_writer_kwargs	)
        return written






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     **a_	)	->int:







        '''simple docstring'''


        _a							=	0
        _a							=	parquet_writer_kwargs.pop("path_or_buf"     ,     a_	)
        _a							=	self.dataset.features.arrow_schema

        _a							=	pq.ParquetWriter(a_     ,     schema=a_     ,     **a_	)

        for offset in logging.tqdm(
            range(0     ,     len(self.dataset	)     ,     a_	)     ,     unit="ba"     ,     disable=not logging.is_progress_bar_enabled()     ,     desc="Creating parquet from Arrow format"     ,     ):
          _a							=	query_table(
              table=self.dataset._data     ,     key=slice(a_     ,     offset + batch_size	)     ,     indices=self.dataset._indices if self.dataset._indices is not None else None     ,     )
          writer.write_table(a_	)
          written += batch.nbytes
        writer.close()
        return written


