"""simple docstring"""




import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import DetrConfig, MaskFormerConfig, SwinConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch

				from transformers import MaskFormerForInstanceSegmentation, MaskFormerModel

				if is_vision_available():
								from transformers import MaskFormerImageProcessor

if is_vision_available():
				from PIL import Image
class     UpperCamelCase__:





						def __init__(       self		,__UpperCAmelCase		,__UpperCAmelCase=2		,__UpperCAmelCase=True		,__UpperCAmelCase=False		,__UpperCAmelCase=10		,__UpperCAmelCase=3		,__UpperCAmelCase=32 * 4		,__UpperCAmelCase=32 * 6		,__UpperCAmelCase=4		,__UpperCAmelCase=32		,)					->				Tuple:
													A__	     =							parent
													A__	     =							batch_size
													A__	     =							is_training
													A__	     =							use_auxiliary_loss
													A__	     =							num_queries
													A__	     =							num_channels
													A__	     =							min_size
													A__	     =							max_size
													A__	     =							num_labels
													A__	     =							mask_feature_size





						def 				snake_case__	(       self     )					->				Dict:
													A__	     =							floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]     ).to(
													    __UpperCAmelCase     )

													A__	     =							torch.ones([self.batch_size, self.min_size, self.max_size]		,device=__UpperCAmelCase     )

													A__	     =							(
													    torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]		,device=__UpperCAmelCase     ) > 0.5
													).float()
													A__	     =							(torch.rand((self.batch_size, self.num_labels)		,device=__UpperCAmelCase     ) > 0.5).long()

													A__	     =							self.get_config()
													return config, pixel_values, pixel_mask, mask_labels, class_labels





						def 				snake_case__	(       self     )					->				Any:
													return MaskFormerConfig.from_backbone_and_decoder_configs(
													    backbone_config=SwinConfig(
													        depths=[1, 1, 1, 1]		,)		,decoder_config=DetrConfig(
													        decoder_ffn_dim=1_28		,num_queries=self.num_queries		,decoder_attention_heads=2		,d_model=self.mask_feature_size		,)		,mask_feature_size=self.mask_feature_size		,fpn_feature_size=self.mask_feature_size		,num_channels=self.num_channels		,num_labels=self.num_labels		,)





						def 				snake_case__	(       self     )					->				int:
													A__     ,  A__     ,  A__     ,  A__     ,  A__	     =							self.prepare_config_and_inputs()
													A__	     =							{'pixel_values': pixel_values, 'pixel_mask': pixel_mask}
													return config, inputs_dict





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase     )					->				Tuple:
													A__	     =							output.encoder_hidden_states
													A__	     =							output.pixel_decoder_hidden_states
													A__	     =							output.transformer_decoder_hidden_states

													self.parent.assertTrue(len(__UpperCAmelCase     )		,len(config.backbone_config.depths     )     )
													self.parent.assertTrue(len(__UpperCAmelCase     )		,len(config.backbone_config.depths     )     )
													self.parent.assertTrue(len(__UpperCAmelCase     )		,config.decoder_config.decoder_layers     )





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase=False     )					->				List[str]:
													with torch.no_grad():
																				A__	     =							MaskFormerModel(config=__UpperCAmelCase     )
																				model.to(__UpperCAmelCase     )
																				model.eval()

																				A__	     =							model(pixel_values=__UpperCAmelCase		,pixel_mask=__UpperCAmelCase     )
																				A__	     =							model(__UpperCAmelCase		,output_hidden_states=__UpperCAmelCase     )
													# the correct shape of output.transformer_decoder_hidden_states ensure the correcteness of the
													# encoder and pixel decoder
													self.parent.assertEqual(
													    output.transformer_decoder_last_hidden_state.shape		,(self.batch_size, self.num_queries, self.mask_feature_size)		,)
													# let's ensure the other two hidden state exists
													self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None     )
													self.parent.assertTrue(output.encoder_last_hidden_state is not None     )

													if output_hidden_states:
																				self.check_output_hidden_state(__UpperCAmelCase		,__UpperCAmelCase     )





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase     )					->				Optional[Any]:
													A__	     =							MaskFormerForInstanceSegmentation(config=__UpperCAmelCase     )
													model.to(__UpperCAmelCase     )
													model.eval()

													def comm_check_on_output(__UpperCAmelCase     ):
																				# let's still check that all the required stuff is there
																				self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None     )
																				self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None     )
																				self.parent.assertTrue(result.encoder_last_hidden_state is not None     )
																				# okay, now we need to check the logits shape
																				# due to the encoder compression, masks have a //4 spatial size
																				self.parent.assertEqual(
																				    result.masks_queries_logits.shape		,(self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)		,)
																				# + 1 for null class
																				self.parent.assertEqual(
																				    result.class_queries_logits.shape		,(self.batch_size, self.num_queries, self.num_labels + 1)     )

													with torch.no_grad():
																				A__	     =							model(pixel_values=__UpperCAmelCase		,pixel_mask=__UpperCAmelCase     )
																				A__	     =							model(__UpperCAmelCase     )

																				comm_check_on_output(__UpperCAmelCase     )

																				A__	     =							model(
																				    pixel_values=__UpperCAmelCase		,pixel_mask=__UpperCAmelCase		,mask_labels=__UpperCAmelCase		,class_labels=__UpperCAmelCase     )

													comm_check_on_output(__UpperCAmelCase     )

													self.parent.assertTrue(result.loss is not None     )
													self.parent.assertEqual(result.loss.shape		,torch.Size([1]     )     )






@require_torch
class     UpperCamelCase__(			__A		,						__A		,						unittest.TestCase ):
						lowerCAmelCase__  :						Tuple													=		(MaskFormerModel, MaskFormerForInstanceSegmentation) if is_torch_available() else ()
						lowerCAmelCase__  :						Any													=		(
						    {'feature-extraction': MaskFormerModel, 'image-segmentation': MaskFormerForInstanceSegmentation}
						    if is_torch_available()
						    else {}
						)

						lowerCAmelCase__  :						Union[str, Any]													=		False
						lowerCAmelCase__  :						Optional[int]													=		False
						lowerCAmelCase__  :						Union[str, Any]													=		False
						lowerCAmelCase__  :						Tuple													=		False





						def 				snake_case__	(       self     )					->				Union[str, Any]:
													A__	     =							MaskFormerModelTester(self     )
													A__	     =							ConfigTester(self		,config_class=__UpperCAmelCase		,has_text_modality=__UpperCAmelCase     )





						def 				snake_case__	(       self     )					->				Optional[Any]:
													self.config_tester.run_common_tests()





						def 				snake_case__	(       self     )					->				Union[str, Any]:
													A__     ,  A__	     =							self.model_tester.prepare_config_and_inputs_for_common()
													self.model_tester.create_and_check_maskformer_model(__UpperCAmelCase		,**__UpperCAmelCase		,output_hidden_states=__UpperCAmelCase     )





						def 				snake_case__	(       self     )					->				Optional[Any]:
													A__	     =							self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_maskformer_instance_segmentation_head_model(*__UpperCAmelCase     )





						@unittest.skip(reason='MaskFormer does not use inputs_embeds'     )
						def 				snake_case__	(       self     )					->				Optional[int]:
													pass





						@unittest.skip(reason='MaskFormer does not have a get_input_embeddings method'     )
						def 				snake_case__	(       self     )					->				Optional[Any]:
													pass





						@unittest.skip(reason='MaskFormer is not a generative model'     )
						def 				snake_case__	(       self     )					->				Optional[Any]:
													pass





						@unittest.skip(reason='MaskFormer does not use token embeddings'     )
						def 				snake_case__	(       self     )					->				Any:
													pass





						@require_torch_multi_gpu
						@unittest.skip(
						    reason='MaskFormer has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'     )
						def 				snake_case__	(       self     )					->				Any:
													pass





						@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'     )
						def 				snake_case__	(       self     )					->				Optional[Any]:
													pass





						def 				snake_case__	(       self     )					->				Dict:
													A__     ,  A__	     =							self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				A__	     =							model_class(__UpperCAmelCase     )
																				A__	     =							inspect.signature(model.forward     )
																				# signature.parameters is an OrderedDict => so arg_names order is deterministic
																				A__	     =							[*signature.parameters.keys()]

																				A__	     =							['pixel_values']
																				self.assertListEqual(arg_names[:1]		,__UpperCAmelCase     )





						@slow
						def 				snake_case__	(       self     )					->				Tuple:
													for model_name in ["facebook/maskformer-swin-small-coco"]:
																				A__	     =							MaskFormerModel.from_pretrained(__UpperCAmelCase     )
																				self.assertIsNotNone(__UpperCAmelCase     )





						def 				snake_case__	(       self     )					->				str:
													A__	     =							(self.model_tester.min_size,) * 2
													A__	     =							{
													    'pixel_values': torch.randn((2, 3, *size)		,device=__UpperCAmelCase     ),
													    'mask_labels': torch.randn((2, 10, *size)		,device=__UpperCAmelCase     ),
													    'class_labels': torch.zeros(2		,10		,device=__UpperCAmelCase     ).long(),
													}

													A__	     =							MaskFormerForInstanceSegmentation(MaskFormerConfig()     ).to(__UpperCAmelCase     )
													A__	     =							model(**__UpperCAmelCase     )
													self.assertTrue(outputs.loss is not None     )





						def 				snake_case__	(       self     )					->				List[str]:
													A__     ,  A__	     =							self.model_tester.prepare_config_and_inputs_for_common()
													self.model_tester.create_and_check_maskformer_model(__UpperCAmelCase		,**__UpperCAmelCase		,output_hidden_states=__UpperCAmelCase     )





						def 				snake_case__	(       self     )					->				Any:
													A__     ,  A__	     =							self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				A__	     =							model_class(__UpperCAmelCase     ).to(__UpperCAmelCase     )
																				A__	     =							model(**__UpperCAmelCase		,output_attentions=__UpperCAmelCase     )
																				self.assertTrue(outputs.attentions is not None     )





						def 				snake_case__	(       self     )					->				int:
													if not self.model_tester.is_training:
																				return
													# only MaskFormerForInstanceSegmentation has the loss
													A__	     =							self.all_model_classes[1]
													A__     ,  A__     ,  A__     ,  A__     ,  A__	     =							self.model_tester.prepare_config_and_inputs()

													A__	     =							model_class(__UpperCAmelCase     )
													model.to(__UpperCAmelCase     )
													model.train()

													A__	     =							model(__UpperCAmelCase		,mask_labels=__UpperCAmelCase		,class_labels=__UpperCAmelCase     ).loss
													loss.backward()





						def 				snake_case__	(       self     )					->				Optional[int]:
													# only MaskFormerForInstanceSegmentation has the loss
													A__	     =							self.all_model_classes[1]
													A__     ,  A__     ,  A__     ,  A__     ,  A__	     =							self.model_tester.prepare_config_and_inputs()
													A__	     =							True
													A__	     =							True

													A__	     =							model_class(__UpperCAmelCase     )
													model.to(__UpperCAmelCase     )
													model.train()

													A__	     =							model(__UpperCAmelCase		,mask_labels=__UpperCAmelCase		,class_labels=__UpperCAmelCase     )

													A__	     =							outputs.encoder_hidden_states[0]
													encoder_hidden_states.retain_grad()

													A__	     =							outputs.pixel_decoder_hidden_states[0]
													pixel_decoder_hidden_states.retain_grad()
													# we requires_grad=True in inputs_embeds (line 2152), the original implementation don't
													A__	     =							outputs.transformer_decoder_hidden_states[0]
													transformer_decoder_hidden_states.retain_grad()

													A__	     =							outputs.attentions[0]
													attentions.retain_grad()

													outputs.loss.backward(retain_graph=__UpperCAmelCase     )

													self.assertIsNotNone(encoder_hidden_states.grad     )
													self.assertIsNotNone(pixel_decoder_hidden_states.grad     )
													self.assertIsNotNone(transformer_decoder_hidden_states.grad     )
													self.assertIsNotNone(attentions.grad     )






__lowerCamelCase											=  1E-4






def   UpperCAmelCase					(	):






							"""simple docstring"""





							A__	     =							Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'				)
							return image


@require_vision
@slow
class     UpperCamelCase__(			unittest.TestCase ):





						@cached_property
						def 				snake_case__	(       self     )					->				Optional[int]:
													return (
													    MaskFormerImageProcessor.from_pretrained('facebook/maskformer-swin-small-coco'     )
													    if is_vision_available()
													    else None
													)





						def 				snake_case__	(       self     )					->				Optional[Any]:
													A__	     =							MaskFormerModel.from_pretrained('facebook/maskformer-swin-small-coco'     ).to(__UpperCAmelCase     )
													A__	     =							self.default_image_processor
													A__	     =							prepare_img()
													A__	     =							image_processor(__UpperCAmelCase		,return_tensors='pt'     ).to(__UpperCAmelCase     )
													A__	     =							inputs['pixel_values'].shape
													# check size is divisible by 32
													self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0     )
													# check size
													self.assertEqual(__UpperCAmelCase		,(1, 3, 8_00, 10_88)     )

													with torch.no_grad():
																				A__	     =							model(**__UpperCAmelCase     )

													A__	     =							torch.tensor(
													    [[-0.0_4_8_2, 0.9_2_2_8, 0.4_9_5_1], [-0.2_5_4_7, 0.8_0_1_7, 0.8_5_2_7], [-0.0_0_6_9, 0.3_3_8_5, -0.0_0_8_9]]     ).to(__UpperCAmelCase     )
													self.assertTrue(
													    torch.allclose(
													        outputs.encoder_last_hidden_state[0, 0, :3, :3]		,__UpperCAmelCase		,atol=__UpperCAmelCase     )     )

													A__	     =							torch.tensor(
													    [[-0.8_4_2_2, -0.8_4_3_4, -0.9_7_1_8], [-1.0_1_4_4, -0.5_5_6_5, -0.4_1_9_5], [-1.0_0_3_8, -0.4_4_8_4, -0.1_9_6_1]]     ).to(__UpperCAmelCase     )
													self.assertTrue(
													    torch.allclose(
													        outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]		,__UpperCAmelCase		,atol=__UpperCAmelCase     )     )

													A__	     =							torch.tensor(
													    [[0.2_8_5_2, -0.0_1_5_9, 0.9_7_3_5], [0.6_2_5_4, 0.1_8_5_8, 0.8_5_2_9], [-0.0_6_8_0, -0.4_1_1_6, 1.8_4_1_3]]     ).to(__UpperCAmelCase     )
													self.assertTrue(
													    torch.allclose(
													        outputs.transformer_decoder_last_hidden_state[0, :3, :3]		,__UpperCAmelCase		,atol=__UpperCAmelCase     )     )





						def 				snake_case__	(       self     )					->				int:
													A__	     =							(
													    MaskFormerForInstanceSegmentation.from_pretrained('facebook/maskformer-swin-small-coco'     )
													    .to(__UpperCAmelCase     )
													    .eval()
													)
													A__	     =							self.default_image_processor
													A__	     =							prepare_img()
													A__	     =							image_processor(__UpperCAmelCase		,return_tensors='pt'     ).to(__UpperCAmelCase     )
													A__	     =							inputs['pixel_values'].shape
													# check size is divisible by 32
													self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0     )
													# check size
													self.assertEqual(__UpperCAmelCase		,(1, 3, 8_00, 10_88)     )

													with torch.no_grad():
																				A__	     =							model(**__UpperCAmelCase     )
													# masks_queries_logits
													A__	     =							outputs.masks_queries_logits
													self.assertEqual(
													    masks_queries_logits.shape		,(1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)		,)
													A__	     =							[
													    [-1.3_7_3_7_1_2_4, -1.7_7_2_4_9_3_7, -1.9_3_6_4_2_3_3],
													    [-1.5_9_7_7_2_8_1, -1.9_8_6_7_9_3_9, -2.1_5_2_3_6_9_5],
													    [-1.5_7_9_5_3_9_8, -1.9_2_6_9_8_3_2, -2.0_9_3_9_4_2],
													]
													A__	     =							torch.tensor(__UpperCAmelCase     ).to(__UpperCAmelCase     )
													self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]		,__UpperCAmelCase		,atol=__UpperCAmelCase     )     )
													# class_queries_logits
													A__	     =							outputs.class_queries_logits
													self.assertEqual(
													    class_queries_logits.shape		,(1, model.config.decoder_config.num_queries, model.config.num_labels + 1)     )
													A__	     =							torch.tensor(
													    [
													        [1.6_512e00, -5.2_572e00, -3.3_519e00],
													        [3.6_169e-02, -5.9_025e00, -2.9_313e00],
													        [1.0_766e-04, -7.7_630e00, -5.1_263e00],
													    ]     ).to(__UpperCAmelCase     )
													self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]		,__UpperCAmelCase		,atol=__UpperCAmelCase     )     )





						def 				snake_case__	(       self     )					->				Optional[Any]:
													A__	     =							(
													    MaskFormerForInstanceSegmentation.from_pretrained('facebook/maskformer-resnet101-coco-stuff'     )
													    .to(__UpperCAmelCase     )
													    .eval()
													)
													A__	     =							self.default_image_processor
													A__	     =							prepare_img()
													A__	     =							image_processor(__UpperCAmelCase		,return_tensors='pt'     ).to(__UpperCAmelCase     )
													A__	     =							inputs['pixel_values'].shape
													# check size is divisible by 32
													self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0     )
													# check size
													self.assertEqual(__UpperCAmelCase		,(1, 3, 8_00, 10_88)     )

													with torch.no_grad():
																				A__	     =							model(**__UpperCAmelCase     )
													# masks_queries_logits
													A__	     =							outputs.masks_queries_logits
													self.assertEqual(
													    masks_queries_logits.shape		,(1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)		,)
													A__	     =							[[-0.9_0_4_6, -2.6_3_6_6, -4.6_0_6_2], [-3.4_1_7_9, -5.7_8_9_0, -8.8_0_5_7], [-4.9_1_7_9, -7.6_5_6_0, -1_0.7_7_1_1]]
													A__	     =							torch.tensor(__UpperCAmelCase     ).to(__UpperCAmelCase     )
													self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]		,__UpperCAmelCase		,atol=__UpperCAmelCase     )     )
													# class_queries_logits
													A__	     =							outputs.class_queries_logits
													self.assertEqual(
													    class_queries_logits.shape		,(1, model.config.decoder_config.num_queries, model.config.num_labels + 1)     )
													A__	     =							torch.tensor(
													    [[4.7_1_8_8, -3.2_5_8_5, -2.8_8_5_7], [6.6_8_7_1, -2.9_1_8_1, -1.2_4_8_7], [7.2_4_4_9, -2.2_7_6_4, -2.1_8_7_4]]     ).to(__UpperCAmelCase     )
													self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]		,__UpperCAmelCase		,atol=__UpperCAmelCase     )     )





						def 				snake_case__	(       self     )					->				Optional[int]:
													A__	     =							(
													    MaskFormerForInstanceSegmentation.from_pretrained('facebook/maskformer-swin-small-coco'     )
													    .to(__UpperCAmelCase     )
													    .eval()
													)
													A__	     =							self.default_image_processor

													A__	     =							image_processor(
													    [np.zeros((3, 8_00, 13_33)     ), np.zeros((3, 8_00, 13_33)     )]		,segmentation_maps=[np.zeros((3_84, 3_84)     ).astype(np.floataa     ), np.zeros((3_84, 3_84)     ).astype(np.floataa     )]		,return_tensors='pt'		,)

													A__	     =							inputs['pixel_values'].to(__UpperCAmelCase     )
													A__	     =							[el.to(__UpperCAmelCase     ) for el in inputs['mask_labels']]
													A__	     =							[el.to(__UpperCAmelCase     ) for el in inputs['class_labels']]

													with torch.no_grad():
																				A__	     =							model(**__UpperCAmelCase     )

													self.assertTrue(outputs.loss is not None     )






"""simple docstring"""




import unittest

import numpy as np

from diffusers import OnnxStableDiffusionInpaintPipelineLegacy
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    load_numpy,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)


if is_onnx_available():
				import onnxruntime as ort
@nightly
@require_onnxruntime
@require_torch_gpu
class     UpperCamelCase__(			unittest.TestCase ):





						@property
						def 				snake_case__	(       self     )					->				int:
													return (
													    "CUDAExecutionProvider",
													    {
													        "gpu_mem_limit": "15000000000",  # 15GB
													        "arena_extend_strategy": "kSameAsRequested",
													    },
													)





						@property
						def 				snake_case__	(       self     )					->				Union[str, Any]:
													A__	     =							ort.SessionOptions()
													A__	     =							False
													return options





						def 				snake_case__	(       self     )					->				str:
													A__	     =							load_image(
													    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
													    '/in_paint/overture-creations-5sI6fQgYIuo.png'     )
													A__	     =							load_image(
													    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
													    '/in_paint/overture-creations-5sI6fQgYIuo_mask.png'     )
													A__	     =							load_numpy(
													    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
													    '/in_paint/red_cat_sitting_on_a_park_bench_onnx.npy'     )

													# using the PNDM scheduler by default
													A__	     =							OnnxStableDiffusionInpaintPipelineLegacy.from_pretrained(
													    'CompVis/stable-diffusion-v1-4'		,revision='onnx'		,safety_checker=__UpperCAmelCase		,feature_extractor=__UpperCAmelCase		,provider=self.gpu_provider		,sess_options=self.gpu_options		,)
													pipe.set_progress_bar_config(disable=__UpperCAmelCase     )

													A__	     =							'A red cat sitting on a park bench'

													A__	     =							np.random.RandomState(0     )
													A__	     =							pipe(
													    prompt=__UpperCAmelCase		,image=__UpperCAmelCase		,mask_image=__UpperCAmelCase		,strength=0.7_5		,guidance_scale=7.5		,num_inference_steps=15		,generator=__UpperCAmelCase		,output_type='np'		,)

													A__	     =							output.images[0]

													assert image.shape == (5_12, 5_12, 3)
													assert np.abs(expected_image - image     ).max() < 1e-2





