



from typing import Any, Dict, List, Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, ChunkPipeline


if is_vision_available():
						from PIL import Image

						from ..image_utils import load_image

if is_torch_available():
						import torch

						from transformers.modeling_outputs import BaseModelOutput

						from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING

UpperCAmelCase_							:      List[Any]      =  logging.get_logger(__name__)






@add_end_docstrings(_UpperCAmelCase     )
class      UpperCamelCase    (						_UpperCAmelCase     ):



	def __init__(       self			,	**UpperCAmelCase__		):
							super().__init__(**UpperCAmelCase__		)

							if self.framework == "tf":
													raise ValueError(F"""The {self.__class__} is only available in PyTorch."""		)

							requires_backends(self			,	"vision"		)
							self.check_model_type(UpperCAmelCase__		)



	def __call__(       self			,	UpperCAmelCase__			,	UpperCAmelCase__ = None			,	**UpperCAmelCase__			,	):
							if "text_queries" in kwargs:
													A__							=       kwargs.pop("text_queries"		)

							if isinstance(UpperCAmelCase__			,	(str, Image.Image)		):
													A__							=       {"image": image, "candidate_labels": candidate_labels}
							else:
													A__							=       image
							A__							=       super().__call__(UpperCAmelCase__			,	**UpperCAmelCase__		)
							return results



	def 			__A							(       self			,	**UpperCAmelCase__		):
							A__							=       {}
							if "threshold" in kwargs:
													A__							=       kwargs["threshold"]
							if "top_k" in kwargs:
													A__							=       kwargs["top_k"]
							return {}, {}, postprocess_params



	def 			__A							(       self			,	UpperCAmelCase__		):
							A__							=       load_image(inputs["image"]		)
							A__							=       inputs["candidate_labels"]
							if isinstance(UpperCAmelCase__			,	UpperCAmelCase__		):
													A__							=       candidate_labels.split(","		)

							A__							=       torch.tensor([[image.height, image.width]]			,	dtype=torch.intaa		)
							for i, candidate_label in enumerate(UpperCAmelCase__		):
													A__							=       self.tokenizer(UpperCAmelCase__			,	return_tensors=self.framework		)
													A__							=       self.image_processor(UpperCAmelCase__			,	return_tensors=self.framework		)
													yield {
													    "is_last": i == len(UpperCAmelCase__		) - 1,
													    "target_size": target_size,
													    "candidate_label": candidate_label,
													    **text_inputs,
													    **image_features,
													}



	def 			__A							(       self			,	UpperCAmelCase__		):
							A__							=       model_inputs.pop("target_size"		)
							A__							=       model_inputs.pop("candidate_label"		)
							A__							=       model_inputs.pop("is_last"		)

							A__							=       self.model(**UpperCAmelCase__		)

							A__							=       {"target_size": target_size, "candidate_label": candidate_label, "is_last": is_last, **outputs}
							return model_outputs



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__=0.1			,	UpperCAmelCase__=None		):
							A__							=       []
							for model_output in model_outputs:
													A__							=       model_output["candidate_label"]
													A__							=       BaseModelOutput(UpperCAmelCase__		)
													A__							=       self.image_processor.post_process_object_detection(
													    outputs=UpperCAmelCase__			,	threshold=UpperCAmelCase__			,	target_sizes=model_output["target_size"]		)[0]

													for index in outputs["scores"].nonzero():
																			A__							=       outputs["scores"][index].item()
																			A__							=       self._get_bounding_box(outputs["boxes"][index][0]		)

																			A__							=       {"score": score, "label": label, "box": box}
																			results.append(UpperCAmelCase__		)

							A__							=       sorted(UpperCAmelCase__			,	key=lambda UpperCAmelCase__		: x["score"]			,	reverse=UpperCAmelCase__		)
							if top_k:
													A__							=       results[:top_k]

							return results



	def 			__A							(       self			,	UpperCAmelCase__		):
							if self.framework != "pt":
													raise ValueError("The ZeroShotObjectDetectionPipeline is only available in PyTorch."		)
							A__						,       A__						,       A__						,       A__							=       box.int().tolist()
							A__							=       {
							    "xmin": xmin,
							    "ymin": ymin,
							    "xmax": xmax,
							    "ymax": ymax,
							}
							return bbox










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    (						_UpperCAmelCase     ):



	def __init__(       self			,	UpperCAmelCase__			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = False			,	UpperCAmelCase__ = False			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = None			,	**UpperCAmelCase__			,	):
							super().__init__(
							    UpperCAmelCase__			,	split=UpperCAmelCase__			,	features=UpperCAmelCase__			,	cache_dir=UpperCAmelCase__			,	keep_in_memory=UpperCAmelCase__			,	streaming=UpperCAmelCase__			,	num_proc=UpperCAmelCase__			,	**UpperCAmelCase__			,	)
							A__							=       field
							A__							=       path_or_paths if isinstance(UpperCAmelCase__			,	UpperCAmelCase__		) else {self.split: path_or_paths}
							A__							=       Json(
							    cache_dir=UpperCAmelCase__			,	data_files=UpperCAmelCase__			,	features=UpperCAmelCase__			,	field=UpperCAmelCase__			,	**UpperCAmelCase__			,	)



	def 			__A							(       self		):
							# Build iterable dataset
							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=UpperCAmelCase__			,	download_mode=UpperCAmelCase__			,	verification_mode=UpperCAmelCase__			,	base_path=UpperCAmelCase__			,	num_proc=self.num_proc			,	)
													A__							=       self.builder.as_dataset(
													    split=self.split			,	verification_mode=UpperCAmelCase__			,	in_memory=self.keep_in_memory		)
							return dataset






class      UpperCamelCase    :



	def __init__(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = None			,	**UpperCAmelCase__			,	):
							if num_proc is not None and num_proc <= 0:
													raise ValueError(F"""num_proc {num_proc} must be an integer > 0."""		)

							A__							=       dataset
							A__							=       path_or_buf
							A__							=       batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
							A__							=       num_proc
							A__							=       "utf-8"
							A__							=       to_json_kwargs



	def 			__A							(       self		):
							A__							=       self.to_json_kwargs.pop("path_or_buf"			,	UpperCAmelCase__		)
							A__							=       self.to_json_kwargs.pop("orient"			,	"records"		)
							A__							=       self.to_json_kwargs.pop("lines"			,	True if orient == "records" else False		)
							A__							=       self.to_json_kwargs.pop("index"			,	False if orient in ["split", "table"] else True		)
							A__							=       self.to_json_kwargs.pop("compression"			,	UpperCAmelCase__		)

							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=UpperCAmelCase__		) as buffer:
																			A__							=       self._write(file_obj=UpperCAmelCase__			,	orient=UpperCAmelCase__			,	lines=UpperCAmelCase__			,	index=UpperCAmelCase__			,	**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."		)
													A__							=       self._write(
													    file_obj=self.path_or_buf			,	orient=UpperCAmelCase__			,	lines=UpperCAmelCase__			,	index=UpperCAmelCase__			,	**self.to_json_kwargs		)
							return written



	def 			__A							(       self			,	UpperCAmelCase__		):
							A__						,       A__						,       A__						,       A__						,       A__							=       args

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



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__			,	):
							A__							=       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"			,	):
																			A__							=       self._batch_json((offset, orient, lines, index, to_json_kwargs)		)
																			written += file_obj.write(UpperCAmelCase__		)
							else:
													A__						,       A__							=       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			,	UpperCAmelCase__			,	UpperCAmelCase__		)]			,	)			,	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(UpperCAmelCase__		)

							return written





