


import unittest

from transformers import is_vision_available
from transformers.pipelines import pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
   from PIL import Image
else:






   class      UpperCamelCase					:


      '''simple docstring'''







      @staticmethod
      def 		UpperCamelCase  (     *UpperCamelCase_   ,							**UpperCamelCase_		):
           pass






@is_pipeline_test
@require_vision
class      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   @require_torch
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[str]     =       pipeline(
            model='''hf-internal-testing/tiny-random-clip-zero-shot-image-classification'''   ,							)
        lowercase_       :Any     =       Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''		)
        lowercase_       :Tuple     =       image_classifier(UpperCamelCase_   ,							candidate_labels=['''a''', '''b''', '''c''']		)

        # The floating scores are so close, we enter floating error approximation and the order is not guaranteed across
        # python and torch versions.
        self.assertIn(
            nested_simplify(UpperCamelCase_		)   ,							[
                [{'''score''': 0.333, '''label''': '''a'''}, {'''score''': 0.333, '''label''': '''b'''}, {'''score''': 0.333, '''label''': '''c'''}],
                [{'''score''': 0.333, '''label''': '''a'''}, {'''score''': 0.333, '''label''': '''c'''}, {'''score''': 0.333, '''label''': '''b'''}],
            ]   ,							)

        lowercase_       :int     =       image_classifier([image] * 5   ,							candidate_labels=['''A''', '''B''', '''C''']   ,							batch_size=2		)
        self.assertEqual(
            nested_simplify(UpperCamelCase_		)   ,							[
                [
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                ],
                [
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                ],
                [
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                ],
                [
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                ],
                [
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                ],
            ]   ,							)
   @require_tf
   def 		UpperCamelCase  (     self		):
        lowercase_       :str     =       pipeline(
            model='''hf-internal-testing/tiny-random-clip-zero-shot-image-classification'''   ,							framework='''tf'''		)
        lowercase_       :Optional[int]     =       Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''		)
        lowercase_       :Tuple     =       image_classifier(UpperCamelCase_   ,							candidate_labels=['''a''', '''b''', '''c''']		)

        self.assertEqual(
            nested_simplify(UpperCamelCase_		)   ,							[{'''score''': 0.333, '''label''': '''a'''}, {'''score''': 0.333, '''label''': '''b'''}, {'''score''': 0.333, '''label''': '''c'''}]   ,							)

        lowercase_       :Dict     =       image_classifier([image] * 5   ,							candidate_labels=['''A''', '''B''', '''C''']   ,							batch_size=2		)
        self.assertEqual(
            nested_simplify(UpperCamelCase_		)   ,							[
                [
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                ],
                [
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                ],
                [
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                ],
                [
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                ],
                [
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                    {'''score''': 0.333, '''label''': ANY(UpperCamelCase_		)},
                ],
            ]   ,							)
   @slow
   @require_torch
   def 		UpperCamelCase  (     self		):
        lowercase_       :Any     =       pipeline(
            task='''zero-shot-image-classification'''   ,							model='''openai/clip-vit-base-patch32'''   ,							)
        # This is an image of 2 cats with remotes and no planes
        lowercase_       :Any     =       Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''		)
        lowercase_       :Dict     =       image_classifier(UpperCamelCase_   ,							candidate_labels=['''cat''', '''plane''', '''remote''']		)

        self.assertEqual(
            nested_simplify(UpperCamelCase_		)   ,							[
                {'''score''': 0.511, '''label''': '''remote'''},
                {'''score''': 0.485, '''label''': '''cat'''},
                {'''score''': 0.004, '''label''': '''plane'''},
            ]   ,							)

        lowercase_       :Any     =       image_classifier([image] * 5   ,							candidate_labels=['''cat''', '''plane''', '''remote''']   ,							batch_size=2		)
        self.assertEqual(
            nested_simplify(UpperCamelCase_		)   ,							[
                [
                    {'''score''': 0.511, '''label''': '''remote'''},
                    {'''score''': 0.485, '''label''': '''cat'''},
                    {'''score''': 0.004, '''label''': '''plane'''},
                ],
            ]
            * 5   ,							)




   @slow
   @require_tf
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[str]     =       pipeline(
            task='''zero-shot-image-classification'''   ,							model='''openai/clip-vit-base-patch32'''   ,							framework='''tf'''		)
        # This is an image of 2 cats with remotes and no planes
        lowercase_       :int     =       Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''		)
        lowercase_       :str     =       image_classifier(UpperCamelCase_   ,							candidate_labels=['''cat''', '''plane''', '''remote''']		)
        self.assertEqual(
            nested_simplify(UpperCamelCase_		)   ,							[
                {'''score''': 0.511, '''label''': '''remote'''},
                {'''score''': 0.485, '''label''': '''cat'''},
                {'''score''': 0.004, '''label''': '''plane'''},
            ]   ,							)

        lowercase_       :Optional[int]     =       image_classifier([image] * 5   ,							candidate_labels=['''cat''', '''plane''', '''remote''']   ,							batch_size=2		)
        self.assertEqual(
            nested_simplify(UpperCamelCase_		)   ,							[
                [
                    {'''score''': 0.511, '''label''': '''remote'''},
                    {'''score''': 0.485, '''label''': '''cat'''},
                    {'''score''': 0.004, '''label''': '''plane'''},
                ],
            ]
            * 5   ,							)





import numpy as np
import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModelWithProjection, PreTrainedModel

from ...utils import logging


SCREAMING_SNAKE_CASE       :		int         =				logging.get_logger(__name__)






class      UpperCamelCase					(		lowercase__     ):


   '''simple docstring'''







   lowercase					:	List[str]									=CLIPConfig

   lowercase					:	Tuple									=["""CLIPEncoderLayer"""]
   def __init__(     self   ,							UpperCamelCase_		):
        super().__init__(UpperCamelCase_		)

        lowercase_       :Union[str, Any]     =       CLIPVisionModelWithProjection(config.vision_config		)

        lowercase_       :str     =       nn.Linear(config.vision_config.projection_dim   ,							1		)
        lowercase_       :Optional[int]     =       nn.Linear(config.vision_config.projection_dim   ,							1		)




   @torch.no_grad()
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_=0.5   ,							UpperCamelCase_=0.5		):
        lowercase_       :Optional[int]     =       self.vision_model(UpperCamelCase_		)[0]

        lowercase_       :Optional[int]     =       self.p_head(UpperCamelCase_		)
        lowercase_       :List[Any]     =       nsfw_detected.flatten()
        lowercase_       :List[str]     =       nsfw_detected > p_threshold
        lowercase_       :Optional[int]     =       nsfw_detected.tolist()

        if any(UpperCamelCase_		):
             logger.warning(
                 '''Potential NSFW content was detected in one or more images. A black image will be returned instead.'''
                 ''' Try again with a different prompt and/or seed.'''		)

        for idx, nsfw_detected_ in enumerate(UpperCamelCase_		):
             if nsfw_detected_:
                  lowercase_       :str     =       np.zeros(images[idx].shape		)

        lowercase_       :int     =       self.w_head(UpperCamelCase_		)
        lowercase_       :Union[str, Any]     =       watermark_detected.flatten()
        lowercase_       :int     =       watermark_detected > w_threshold
        lowercase_       :int     =       watermark_detected.tolist()

        if any(UpperCamelCase_		):
             logger.warning(
                 '''Potential watermarked content was detected in one or more images. A black image will be returned instead.'''
                 ''' Try again with a different prompt and/or seed.'''		)

        for idx, watermark_detected_ in enumerate(UpperCamelCase_		):
             if watermark_detected_:
                  lowercase_       :Dict     =       np.zeros(images[idx].shape		)

        return images, nsfw_detected, watermark_detected

