






'''simple docstring'''
from copy import deepcopy
from typing import Optional, Union

import numpy as np

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding
from ...utils import TensorType, is_tf_available, is_torch_available


if is_torch_available():
 import torch

if is_tf_available():
 import tensorflow as tf


class        _UpperCAmelCase						(		__snake_case			):
    __lowerCamelCase:   Tuple		       =    ['image_processor']
    __lowerCamelCase:   List[Any]		       =    'SamImageProcessor'




    def __init__(			self  :       Tuple  ,     a  :       str			):






       '''simple docstring'''






       super().__init__(__UpperCamelCase			)
       lowercase_	:    Optional[int]           =     self.image_processor
       lowercase_	:    Tuple           =     -1_0
       lowercase_	:    Optional[int]           =     self.image_processor.size["longest_edge"]




    def __call__(			self  :       Dict  ,     a  :       List[str]=None  ,     a  :       Tuple=None  ,     a  :       int=None  ,     a  :       Optional[Any]=None  ,     a  :       int = None  ,     **a  :       List[str]  ,     ):






       '''simple docstring'''






       lowercase_	:    List[str]           =     self.image_processor(
           __UpperCamelCase  ,     return_tensors=__UpperCamelCase  ,     **__UpperCamelCase  ,     )

       # pop arguments that are not used in the foward but used nevertheless
       lowercase_	:    Tuple           =     encoding_image_processor["original_sizes"]

       if hasattr(__UpperCamelCase  ,     "numpy"			):  # Checks if Torch or TF tensor
          lowercase_	:    Optional[int]           =     original_sizes.numpy()

       lowercase_    ,		lowercase_    ,		lowercase_	:    int           =     self._check_and_preprocess_points(
           input_points=__UpperCamelCase  ,     input_labels=__UpperCamelCase  ,     input_boxes=__UpperCamelCase  ,     )

       lowercase_	:    int           =     self._normalize_and_convert(
           __UpperCamelCase  ,     __UpperCamelCase  ,     input_points=__UpperCamelCase  ,     input_labels=__UpperCamelCase  ,     input_boxes=__UpperCamelCase  ,     return_tensors=__UpperCamelCase  ,     )

       return encoding_image_processor




    def 				lowerCAmelCase__				(			self  :       Union[str, Any]  ,     a  :       List[str]  ,     a  :       Tuple  ,     a  :       Optional[int]=None  ,     a  :       List[str]=None  ,     a  :       Dict=None  ,     a  :       List[str]="pt"  ,     ):






       '''simple docstring'''






       if input_points is not None:
          if len(__UpperCamelCase			) != len(__UpperCamelCase			):
             lowercase_	:    Optional[int]           =     [
                 self._normalize_coordinates(self.target_size  ,     __UpperCamelCase  ,     original_sizes[0]			) for point in input_points
             ]
          else:
             lowercase_	:    int           =     [
                 self._normalize_coordinates(self.target_size  ,     __UpperCamelCase  ,     __UpperCamelCase			)
                 for point, original_size in zip(__UpperCamelCase  ,     __UpperCamelCase			)
             ]
          # check that all arrays have the same shape
          if not all(point.shape == input_points[0].shape for point in input_points			):
             if input_labels is not None:
                lowercase_    ,		lowercase_	:    Optional[Any]           =     self._pad_points_and_labels(__UpperCamelCase  ,     __UpperCamelCase			)

          lowercase_	:    Dict           =     np.array(__UpperCamelCase			)

       if input_labels is not None:
          lowercase_	:    int           =     np.array(__UpperCamelCase			)

       if input_boxes is not None:
          if len(__UpperCamelCase			) != len(__UpperCamelCase			):
             lowercase_	:    Union[str, Any]           =     [
                 self._normalize_coordinates(self.target_size  ,     __UpperCamelCase  ,     original_sizes[0]  ,     is_bounding_box=__UpperCamelCase			)
                 for box in input_boxes
             ]
          else:
             lowercase_	:    int           =     [
                 self._normalize_coordinates(self.target_size  ,     __UpperCamelCase  ,     __UpperCamelCase  ,     is_bounding_box=__UpperCamelCase			)
                 for box, original_size in zip(__UpperCamelCase  ,     __UpperCamelCase			)
             ]
          lowercase_	:    Dict           =     np.array(__UpperCamelCase			)

       if input_boxes is not None:
          if return_tensors == "pt":
             lowercase_	:    int           =     torch.from_numpy(__UpperCamelCase			)
             # boxes batch size of 1 by default
             lowercase_	:    Optional[Any]           =     input_boxes.unsqueeze(1			) if len(input_boxes.shape			) != 3 else input_boxes
          elif return_tensors == "tf":
             lowercase_	:    Union[str, Any]           =     tf.convert_to_tensor(__UpperCamelCase			)
             # boxes batch size of 1 by default
             lowercase_	:    Optional[Any]           =     tf.expand_dims(__UpperCamelCase  ,     1			) if len(input_boxes.shape			) != 3 else input_boxes
          encoding_image_processor.update({"input_boxes": input_boxes}			)
       if input_points is not None:
          if return_tensors == "pt":
             lowercase_	:    str           =     torch.from_numpy(__UpperCamelCase			)
             # point batch size of 1 by default
             lowercase_	:    List[str]           =     input_points.unsqueeze(1			) if len(input_points.shape			) != 4 else input_points
          elif return_tensors == "tf":
             lowercase_	:    List[Any]           =     tf.convert_to_tensor(__UpperCamelCase			)
             # point batch size of 1 by default
             lowercase_	:    Union[str, Any]           =     tf.expand_dims(__UpperCamelCase  ,     1			) if len(input_points.shape			) != 4 else input_points
          encoding_image_processor.update({"input_points": input_points}			)
       if input_labels is not None:
          if return_tensors == "pt":
             lowercase_	:    List[str]           =     torch.from_numpy(__UpperCamelCase			)
             # point batch size of 1 by default
             lowercase_	:    Tuple           =     input_labels.unsqueeze(1			) if len(input_labels.shape			) != 3 else input_labels
          elif return_tensors == "tf":
             lowercase_	:    str           =     tf.convert_to_tensor(__UpperCamelCase			)
             # point batch size of 1 by default
             lowercase_	:    Dict           =     tf.expand_dims(__UpperCamelCase  ,     1			) if len(input_labels.shape			) != 3 else input_labels
          encoding_image_processor.update({"input_labels": input_labels}			)

       return encoding_image_processor




    def 				lowerCAmelCase__				(			self  :       Optional[int]  ,     a  :       Optional[int]  ,     a  :       Optional[int]			):






       '''simple docstring'''






       lowercase_	:    List[Any]           =     max([point.shape[0] for point in input_points]			)
       lowercase_	:    List[Any]           =     []
       for i, point in enumerate(__UpperCamelCase			):
          if point.shape[0] != expected_nb_points:
             lowercase_	:    Any           =     np.concatenate(
                 [point, np.zeros((expected_nb_points - point.shape[0], 2)			) + self.point_pad_value]  ,     axis=0			)
             lowercase_	:    Optional[int]           =     np.append(input_labels[i]  ,     [self.point_pad_value]			)
          processed_input_points.append(__UpperCamelCase			)
       lowercase_	:    Union[str, Any]           =     processed_input_points
       return input_points, input_labels




    def 				lowerCAmelCase__				(			self  :       Optional[int]  ,     a  :       Optional[int]  ,     a  :       Tuple  ,     a  :       Optional[int]  ,     a  :       str=False			):






       '''simple docstring'''






       lowercase_    ,		lowercase_	:    int           =     original_size
       lowercase_    ,		lowercase_	:    Optional[int]           =     self.image_processor._get_preprocess_shape(__UpperCamelCase  ,     longest_edge=__UpperCamelCase			)
       lowercase_	:    Any           =     deepcopy(__UpperCamelCase			).astype(__UpperCamelCase			)

       if is_bounding_box:
          lowercase_	:    Any           =     coords.reshape(-1  ,     2  ,     2			)

       lowercase_	:    int           =     coords[..., 0] * (new_w / old_w)
       lowercase_	:    str           =     coords[..., 1] * (new_h / old_h)

       if is_bounding_box:
          lowercase_	:    Optional[int]           =     coords.reshape(-1  ,     4			)

       return coords




    def 				lowerCAmelCase__				(			self  :       Optional[Any]  ,     a  :       str=None  ,     a  :       List[str]=None  ,     a  :       Optional[int]=None  ,     ):






       '''simple docstring'''






       if input_points is not None:
          if hasattr(__UpperCamelCase  ,     "numpy"			):  # Checks for TF or Torch tensor
             lowercase_	:    Optional[int]           =     input_points.numpy().tolist()

          if not isinstance(__UpperCamelCase  ,     __UpperCamelCase			) or not isinstance(input_points[0]  ,     __UpperCamelCase			):
             raise ValueError("Input points must be a list of list of floating points."			)
          lowercase_	:    Optional[int]           =     [np.array(__UpperCamelCase			) for input_point in input_points]
       else:
          lowercase_	:    Optional[int]           =     None

       if input_labels is not None:
          if hasattr(__UpperCamelCase  ,     "numpy"			):
             lowercase_	:    Dict           =     input_labels.numpy().tolist()

          if not isinstance(__UpperCamelCase  ,     __UpperCamelCase			) or not isinstance(input_labels[0]  ,     __UpperCamelCase			):
             raise ValueError("Input labels must be a list of list integers."			)
          lowercase_	:    Dict           =     [np.array(__UpperCamelCase			) for label in input_labels]
       else:
          lowercase_	:    Any           =     None

       if input_boxes is not None:
          if hasattr(__UpperCamelCase  ,     "numpy"			):
             lowercase_	:    Any           =     input_boxes.numpy().tolist()

          if (
              not isinstance(__UpperCamelCase  ,     __UpperCamelCase			)
              or not isinstance(input_boxes[0]  ,     __UpperCamelCase			)
              or not isinstance(input_boxes[0][0]  ,     __UpperCamelCase			)
          ):
             raise ValueError("Input boxes must be a list of list of list of floating points."			)
          lowercase_	:    int           =     [np.array(__UpperCamelCase			).astype(np.floataa			) for box in input_boxes]
       else:
          lowercase_	:    int           =     None

       return input_points, input_labels, input_boxes




    @property
    def 				lowerCAmelCase__				(			self  :       int			):






       '''simple docstring'''






       lowercase_	:    List[Any]           =     self.image_processor.model_input_names
       return list(dict.fromkeys(__UpperCamelCase			)			)





    def 				lowerCAmelCase__				(			self  :       Dict  ,     *a  :       Tuple  ,     **a  :       List[str]			):






       '''simple docstring'''






       return self.image_processor.post_process_masks(*__UpperCamelCase  ,     **__UpperCamelCase			)












'''simple docstring'''
import multiprocessing
import os
from typing import BinaryIO, Optional, Union

import fsspec

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.json.json import Json
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader


class        _UpperCAmelCase						(		snake_case			):




    def __init__(			self  :       Tuple  ,     a  :       NestedDataStructureLike[PathLike]  ,     a  :       Optional[NamedSplit] = None  ,     a  :       Optional[Features] = None  ,     a  :       str = None  ,     a  :       bool = False  ,     a  :       bool = False  ,     a  :       Optional[str] = None  ,     a  :       Optional[int] = None  ,     **a  :       List[Any]  ,     ):






       '''simple docstring'''






       super().__init__(
           a  ,     split=a  ,     features=a  ,     cache_dir=a  ,     keep_in_memory=a  ,     streaming=a  ,     num_proc=a  ,     **a  ,     )
       lowercase_	:    str           =     field
       lowercase_	:    Optional[Any]           =     path_or_paths if isinstance(a  ,     a			) else {self.split: path_or_paths}
       lowercase_	:    Any           =     Json(
           cache_dir=a  ,     data_files=a  ,     features=a  ,     field=a  ,     **a  ,     )





    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       if self.streaming:
          lowercase_	:    Any           =     self.builder.as_streaming_dataset(split=self.split			)
       # Build regular (map-style) dataset
       else:
          lowercase_	:    Dict           =     None
          lowercase_	:    Optional[int]           =     None
          lowercase_	:    str           =     None
          lowercase_	:    str           =     None

          self.builder.download_and_prepare(
              download_config=a  ,     download_mode=a  ,     verification_mode=a  ,     base_path=a  ,     num_proc=self.num_proc  ,     )
          lowercase_	:    int           =     self.builder.as_dataset(
              split=self.split  ,     verification_mode=a  ,     in_memory=self.keep_in_memory			)
       return dataset


class        _UpperCAmelCase						:




    def __init__(			self  :       str  ,     a  :       Dataset  ,     a  :       Union[PathLike, BinaryIO]  ,     a  :       Optional[int] = None  ,     a  :       Optional[int] = None  ,     **a  :       List[Any]  ,     ):






       '''simple docstring'''






       if num_proc is not None and num_proc <= 0:
          raise ValueError(f"""num_proc {num_proc} must be an integer > 0."""			)

       lowercase_	:    Dict           =     dataset
       lowercase_	:    Optional[int]           =     path_or_buf
       lowercase_	:    List[Any]           =     batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
       lowercase_	:    Optional[Any]           =     num_proc
       lowercase_	:    List[Any]           =     "utf-8"
       lowercase_	:    List[str]           =     to_json_kwargs




    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       lowercase_	:    str           =     self.to_json_kwargs.pop("path_or_buf"  ,     a			)
       lowercase_	:    Any           =     self.to_json_kwargs.pop("orient"  ,     "records"			)
       lowercase_	:    Any           =     self.to_json_kwargs.pop("lines"  ,     True if orient == "records" else False			)
       lowercase_	:    List[str]           =     self.to_json_kwargs.pop("index"  ,     False if orient in ["split", "table"] else True			)
       lowercase_	:    int           =     self.to_json_kwargs.pop("compression"  ,     a			)

       if compression not in [None, "infer", "gzip", "bz2", "xz"]:
          raise NotImplementedError(f"""`datasets` currently does not support {compression} compression"""			)

       if isinstance(self.path_or_buf  ,     (str, bytes, os.PathLike)			):
          with fsspec.open(self.path_or_buf  ,     "wb"  ,     compression=a			) as buffer:
             lowercase_	:    Dict           =     self._write(file_obj=a  ,     orient=a  ,     lines=a  ,     index=a  ,     **self.to_json_kwargs			)
       else:
          if compression:
             raise NotImplementedError(
                 f"""The compression parameter is not supported when writing to a buffer, but compression={compression}"""
                 " was passed. Please provide a local path instead."			)
          lowercase_	:    Dict           =     self._write(
              file_obj=self.path_or_buf  ,     orient=a  ,     lines=a  ,     index=a  ,     **self.to_json_kwargs			)
       return written




    def 				lowerCAmelCase__				(			self  :       Optional[int]  ,     a  :       List[str]			):






       '''simple docstring'''






       lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_	:    List[Any]           =     args

       lowercase_	:    Optional[int]           =     query_table(
           table=self.dataset.data  ,     key=slice(a  ,     offset + self.batch_size			)  ,     indices=self.dataset._indices  ,     )
       lowercase_	:    Dict           =     batch.to_pandas().to_json(
           path_or_buf=a  ,     orient=a  ,     lines=a  ,     index=a  ,     **a			)
       if not json_str.endswith("\n"			):
          json_str += "\n"
       return json_str.encode(self.encoding			)





    def 				lowerCAmelCase__				(			self  :       int  ,     a  :       BinaryIO  ,     a  :       int  ,     a  :       str  ,     a  :       Union[str, Any]  ,     **a  :       str  ,     ):






       '''simple docstring'''






       lowercase_	:    Union[str, Any]           =     0

       if self.num_proc is None or self.num_proc == 1:
          for offset in logging.tqdm(
              range(0  ,     len(self.dataset			)  ,     self.batch_size			)  ,     unit="ba"  ,     disable=not logging.is_progress_bar_enabled()  ,     desc="Creating json from Arrow format"  ,     ):
             lowercase_	:    Dict           =     self._batch_json((offset, orient, lines, index, to_json_kwargs)			)
             written += file_obj.write(a			)
       else:
          lowercase_    ,		lowercase_	:    Any           =     len(self.dataset			), self.batch_size
          with multiprocessing.Pool(self.num_proc			) as pool:
             for json_str in logging.tqdm(
                 pool.imap(
                     self._batch_json  ,     [(offset, orient, lines, index, to_json_kwargs) for offset in range(0  ,     a  ,     a			)]  ,     )  ,     total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size  ,     unit="ba"  ,     disable=not logging.is_progress_bar_enabled()  ,     desc="Creating json from Arrow format"  ,     ):
                written += file_obj.write(a			)

       return written




