Spaces:
Running
on
T4
Running
on
T4
| from typing import List, Union | |
| from ..utils import ( | |
| add_end_docstrings, | |
| is_tf_available, | |
| is_torch_available, | |
| is_vision_available, | |
| logging, | |
| requires_backends, | |
| ) | |
| from .base import PIPELINE_INIT_ARGS, Pipeline | |
| if is_vision_available(): | |
| from PIL import Image | |
| from ..image_utils import load_image | |
| if is_tf_available(): | |
| import tensorflow as tf | |
| from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES | |
| from ..tf_utils import stable_softmax | |
| if is_torch_available(): | |
| from ..models.auto.modeling_auto import MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES | |
| logger = logging.get_logger(__name__) | |
| class ImageClassificationPipeline(Pipeline): | |
| """ | |
| Image classification pipeline using any `AutoModelForImageClassification`. This pipeline predicts the class of an | |
| image. | |
| Example: | |
| ```python | |
| >>> from transformers import pipeline | |
| >>> classifier = pipeline(model="microsoft/beit-base-patch16-224-pt22k-ft22k") | |
| >>> classifier("https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png") | |
| [{'score': 0.442, 'label': 'macaw'}, {'score': 0.088, 'label': 'popinjay'}, {'score': 0.075, 'label': 'parrot'}, {'score': 0.073, 'label': 'parodist, lampooner'}, {'score': 0.046, 'label': 'poll, poll_parrot'}] | |
| ``` | |
| Learn more about the basics of using a pipeline in the [pipeline tutorial](../pipeline_tutorial) | |
| This image classification pipeline can currently be loaded from [`pipeline`] using the following task identifier: | |
| `"image-classification"`. | |
| See the list of available models on | |
| [huggingface.co/models](https://huggingface.co/models?filter=image-classification). | |
| """ | |
| def __init__(self, *args, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| requires_backends(self, "vision") | |
| self.check_model_type( | |
| TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES | |
| if self.framework == "tf" | |
| else MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES | |
| ) | |
| def _sanitize_parameters(self, top_k=None, timeout=None): | |
| preprocess_params = {} | |
| if timeout is not None: | |
| preprocess_params["timeout"] = timeout | |
| postprocess_params = {} | |
| if top_k is not None: | |
| postprocess_params["top_k"] = top_k | |
| return preprocess_params, {}, postprocess_params | |
| def __call__(self, images: Union[str, List[str], "Image.Image", List["Image.Image"]], **kwargs): | |
| """ | |
| Assign labels to the image(s) passed as inputs. | |
| Args: | |
| images (`str`, `List[str]`, `PIL.Image` or `List[PIL.Image]`): | |
| The pipeline handles three types of images: | |
| - A string containing a http link pointing to an image | |
| - A string containing a local path to an image | |
| - An image loaded in PIL directly | |
| The pipeline accepts either a single image or a batch of images, which must then be passed as a string. | |
| Images in a batch must all be in the same format: all as http links, all as local paths, or all as PIL | |
| images. | |
| top_k (`int`, *optional*, defaults to 5): | |
| The number of top labels that will be returned by the pipeline. If the provided number is higher than | |
| the number of labels available in the model configuration, it will default to the number of labels. | |
| timeout (`float`, *optional*, defaults to None): | |
| The maximum time in seconds to wait for fetching images from the web. If None, no timeout is set and | |
| the call may block forever. | |
| Return: | |
| A dictionary or a list of dictionaries containing result. If the input is a single image, will return a | |
| dictionary, if the input is a list of several images, will return a list of dictionaries corresponding to | |
| the images. | |
| The dictionaries contain the following keys: | |
| - **label** (`str`) -- The label identified by the model. | |
| - **score** (`int`) -- The score attributed by the model for that label. | |
| """ | |
| return super().__call__(images, **kwargs) | |
| def preprocess(self, image, timeout=None): | |
| image = load_image(image, timeout=timeout) | |
| model_inputs = self.image_processor(images=image, return_tensors=self.framework) | |
| return model_inputs | |
| def _forward(self, model_inputs): | |
| model_outputs = self.model(**model_inputs) | |
| return model_outputs | |
| def postprocess(self, model_outputs, top_k=5): | |
| if top_k > self.model.config.num_labels: | |
| top_k = self.model.config.num_labels | |
| if self.framework == "pt": | |
| probs = model_outputs.logits.softmax(-1)[0] | |
| scores, ids = probs.topk(top_k) | |
| elif self.framework == "tf": | |
| probs = stable_softmax(model_outputs.logits, axis=-1)[0] | |
| topk = tf.math.top_k(probs, k=top_k) | |
| scores, ids = topk.values.numpy(), topk.indices.numpy() | |
| else: | |
| raise ValueError(f"Unsupported framework: {self.framework}") | |
| scores = scores.tolist() | |
| ids = ids.tolist() | |
| return [{"score": score, "label": self.model.config.id2label[_id]} for score, _id in zip(scores, ids)] | |