


'''simple docstring'''
import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextConfig, SegformerImageProcessor, UperNetConfig, UperNetForSemanticSegmentation







def       a_				(	UpperCamelCase_ ):
							A_			    =		3_8_4
							if "tiny" in model_name:
														A_			    =		[3, 3, 9, 3]
														A_			    =		[9_6, 1_9_2, 3_8_4, 7_6_8]
							if "small" in model_name:
														A_			    =		[3, 3, 2_7, 3]
														A_			    =		[9_6, 1_9_2, 3_8_4, 7_6_8]
							if "base" in model_name:
														A_			    =		[3, 3, 2_7, 3]
														A_			    =		[1_2_8, 2_5_6, 5_1_2, 1_0_2_4]
														A_			    =		5_1_2
							if "large" in model_name:
														A_			    =		[3, 3, 2_7, 3]
														A_			    =		[1_9_2, 3_8_4, 7_6_8, 1_5_3_6]
														A_			    =		7_6_8
							if "xlarge" in model_name:
														A_			    =		[3, 3, 2_7, 3]
														A_			    =		[2_5_6, 5_1_2, 1_0_2_4, 2_0_4_8]
														A_			    =		1_0_2_4

							# set label information
							A_			    =		1_5_0
							A_			    =		"huggingface/label-files"
							A_			    =		"ade20k-id2label.json"
							A_			    =		json.load(open(hf_hub_download(UpperCamelCase_						,  UpperCamelCase_						,  repo_type="dataset" )						,  "r" ) )
							A_			    =		{int(UpperCamelCase_ ): v for k, v in idalabel.items()}
							A_			    =		{v: k for k, v in idalabel.items()}

							A_			    =		ConvNextConfig(
							    depths=UpperCamelCase_						,  hidden_sizes=UpperCamelCase_						,  out_features=["stage1", "stage2", "stage3", "stage4"] )
							A_			    =		UperNetConfig(
							    backbone_config=UpperCamelCase_						,  auxiliary_in_channels=UpperCamelCase_						,  num_labels=UpperCamelCase_						,  idalabel=UpperCamelCase_						,  labelaid=UpperCamelCase_						,  )

							return config







def       a_				(	UpperCamelCase_ ):
							A_			    =		[]

							# fmt: off
							# stem
							rename_keys.append(("backbone.downsample_layers.0.0.weight", "backbone.embeddings.patch_embeddings.weight") )
							rename_keys.append(("backbone.downsample_layers.0.0.bias", "backbone.embeddings.patch_embeddings.bias") )
							rename_keys.append(("backbone.downsample_layers.0.1.weight", "backbone.embeddings.layernorm.weight") )
							rename_keys.append(("backbone.downsample_layers.0.1.bias", "backbone.embeddings.layernorm.bias") )
							# stages
							for i in range(len(config.backbone_config.depths ) ):
														for j in range(config.backbone_config.depths[i] ):
																					rename_keys.append((f"backbone.stages.{i}.{j}.gamma", f"backbone.encoder.stages.{i}.layers.{j}.layer_scale_parameter") )
																					rename_keys.append((f"backbone.stages.{i}.{j}.depthwise_conv.weight", f"backbone.encoder.stages.{i}.layers.{j}.dwconv.weight") )
																					rename_keys.append((f"backbone.stages.{i}.{j}.depthwise_conv.bias", f"backbone.encoder.stages.{i}.layers.{j}.dwconv.bias") )
																					rename_keys.append((f"backbone.stages.{i}.{j}.norm.weight", f"backbone.encoder.stages.{i}.layers.{j}.layernorm.weight") )
																					rename_keys.append((f"backbone.stages.{i}.{j}.norm.bias", f"backbone.encoder.stages.{i}.layers.{j}.layernorm.bias") )
																					rename_keys.append((f"backbone.stages.{i}.{j}.pointwise_conv1.weight", f"backbone.encoder.stages.{i}.layers.{j}.pwconv1.weight") )
																					rename_keys.append((f"backbone.stages.{i}.{j}.pointwise_conv1.bias", f"backbone.encoder.stages.{i}.layers.{j}.pwconv1.bias") )
																					rename_keys.append((f"backbone.stages.{i}.{j}.pointwise_conv2.weight", f"backbone.encoder.stages.{i}.layers.{j}.pwconv2.weight") )
																					rename_keys.append((f"backbone.stages.{i}.{j}.pointwise_conv2.bias", f"backbone.encoder.stages.{i}.layers.{j}.pwconv2.bias") )
														if i > 0:
																					rename_keys.append((f"backbone.downsample_layers.{i}.0.weight", f"backbone.encoder.stages.{i}.downsampling_layer.0.weight") )
																					rename_keys.append((f"backbone.downsample_layers.{i}.0.bias", f"backbone.encoder.stages.{i}.downsampling_layer.0.bias") )
																					rename_keys.append((f"backbone.downsample_layers.{i}.1.weight", f"backbone.encoder.stages.{i}.downsampling_layer.1.weight") )
																					rename_keys.append((f"backbone.downsample_layers.{i}.1.bias", f"backbone.encoder.stages.{i}.downsampling_layer.1.bias") )

														rename_keys.append((f"backbone.norm{i}.weight", f"backbone.hidden_states_norms.stage{i+1}.weight") )
														rename_keys.append((f"backbone.norm{i}.bias", f"backbone.hidden_states_norms.stage{i+1}.bias") )

							# decode head
							rename_keys.extend(
							    [
							        ("decode_head.conv_seg.weight", "decode_head.classifier.weight"),
							        ("decode_head.conv_seg.bias", "decode_head.classifier.bias"),
							        ("auxiliary_head.conv_seg.weight", "auxiliary_head.classifier.weight"),
							        ("auxiliary_head.conv_seg.bias", "auxiliary_head.classifier.bias"),
							    ] )
							# fmt: on

							return rename_keys







def       a_				(	UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_ ):
							A_			    =		dct.pop(UpperCamelCase_ )
							A_			    =		val







def       a_				(	UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_ ):
							A_			    =		{
							    "upernet-convnext-tiny": "https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_tiny_fp16_512x512_160k_ade20k/upernet_convnext_tiny_fp16_512x512_160k_ade20k_20220227_124553-cad485de.pth",
							    "upernet-convnext-small": "https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_small_fp16_512x512_160k_ade20k/upernet_convnext_small_fp16_512x512_160k_ade20k_20220227_131208-1b1e394f.pth",
							    "upernet-convnext-base": "https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_base_fp16_512x512_160k_ade20k/upernet_convnext_base_fp16_512x512_160k_ade20k_20220227_181227-02a24fc6.pth",
							    "upernet-convnext-large": "https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_large_fp16_640x640_160k_ade20k/upernet_convnext_large_fp16_640x640_160k_ade20k_20220226_040532-e57aa54d.pth",
							    "upernet-convnext-xlarge": "https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_xlarge_fp16_640x640_160k_ade20k/upernet_convnext_xlarge_fp16_640x640_160k_ade20k_20220226_080344-95fc38c2.pth",
							}
							A_			    =		model_name_to_url[model_name]
							A_			    =		torch.hub.load_state_dict_from_url(UpperCamelCase_						,  map_location="cpu" )["state_dict"]

							A_			    =		get_upernet_config(UpperCamelCase_ )
							A_			    =		UperNetForSemanticSegmentation(UpperCamelCase_ )
							model.eval()

							# replace "bn" => "batch_norm"
							for key in state_dict.copy().keys():
														A_			    =		state_dict.pop(UpperCamelCase_ )
														if "bn" in key:
																					A_			    =		key.replace("bn"						,  "batch_norm" )
														A_			    =		val

							# rename keys
							A_			    =		create_rename_keys(UpperCamelCase_ )
							for src, dest in rename_keys:
														rename_key(UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_ )

							model.load_state_dict(UpperCamelCase_ )

							# verify on image
							A_			    =		"https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"
							A_			    =		Image.open(requests.get(UpperCamelCase_						,  stream=UpperCamelCase_ ).raw ).convert("RGB" )

							A_			    =		SegformerImageProcessor()
							A_			    =		processor(UpperCamelCase_						,  return_tensors="pt" ).pixel_values

							with torch.no_grad():
														A_			    =		model(UpperCamelCase_ )

							if model_name == "upernet-convnext-tiny":
														A_			    =		torch.tensor(
														    [[-8.8110, -8.8110, -8.6521], [-8.8110, -8.8110, -8.6521], [-8.7746, -8.7746, -8.6130]] )
							elif model_name == "upernet-convnext-small":
														A_			    =		torch.tensor(
														    [[-8.8236, -8.8236, -8.6771], [-8.8236, -8.8236, -8.6771], [-8.7638, -8.7638, -8.6240]] )
							elif model_name == "upernet-convnext-base":
														A_			    =		torch.tensor(
														    [[-8.8558, -8.8558, -8.6905], [-8.8558, -8.8558, -8.6905], [-8.7669, -8.7669, -8.6021]] )
							elif model_name == "upernet-convnext-large":
														A_			    =		torch.tensor(
														    [[-8.6660, -8.6660, -8.6210], [-8.6660, -8.6660, -8.6210], [-8.6310, -8.6310, -8.5964]] )
							elif model_name == "upernet-convnext-xlarge":
														A_			    =		torch.tensor(
														    [[-8.4980, -8.4980, -8.3977], [-8.4980, -8.4980, -8.3977], [-8.4379, -8.4379, -8.3412]] )
							print("Logits:"						,  outputs.logits[0, 0, :3, :3] )
							assert torch.allclose(outputs.logits[0, 0, :3, :3]						,  UpperCamelCase_						,  atol=1e-4 )
							print("Looks ok!" )

							if pytorch_dump_folder_path is not None:
														print(f"Saving model {model_name} to {pytorch_dump_folder_path}" )
														model.save_pretrained(UpperCamelCase_ )
														print(f"Saving processor to {pytorch_dump_folder_path}" )
														processor.save_pretrained(UpperCamelCase_ )

							if push_to_hub:
														print(f"Pushing model and processor for {model_name} to hub" )
														model.push_to_hub(f"openmmlab/{model_name}" )
														processor.push_to_hub(f"openmmlab/{model_name}" )


if __name__ == "__main__":
							__SCREAMING_SNAKE_CASE					:				List[Any]									=		argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '''--model_name''',
							    default='''upernet-convnext-tiny''',
							    type=str,
							    choices=[F"upernet-convnext-{size}" for size in ['''tiny''', '''small''', '''base''', '''large''', '''xlarge''']],
							    help='''Name of the ConvNext UperNet model you\'d like to convert.''',
							)
							parser.add_argument(
							    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
							)
							parser.add_argument(
							    '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
							)

							__SCREAMING_SNAKE_CASE					:				Optional[int]									=		parser.parse_args()
							convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









'''simple docstring'''
from typing import Dict, Iterable, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import normalize, rescale, resize, to_channel_dimension_format, to_pil_image
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_pytesseract_available, is_vision_available, logging, requires_backends


if is_vision_available():
							import PIL

# soft dependency
if is_pytesseract_available():
							import pytesseract

__SCREAMING_SNAKE_CASE					:				List[str]									=		logging.get_logger(__name__)







def       a_				(	UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_ ):
							return [
							    int(1_0_0_0 * (box[0] / width) ),
							    int(1_0_0_0 * (box[1] / height) ),
							    int(1_0_0_0 * (box[2] / width) ),
							    int(1_0_0_0 * (box[3] / height) ),
							]







def       a_				(	UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_ ):
							A_			    =		to_pil_image(UpperCamelCase_ )
							A_    ,					A_			    =		pil_image.size
							A_			    =		pytesseract.image_to_data(UpperCamelCase_						,  lang=UpperCamelCase_						,  output_type="dict"						,  config=UpperCamelCase_ )
							A_    ,					A_    ,					A_    ,					A_    ,					A_			    =		data["text"], data["left"], data["top"], data["width"], data["height"]

							# filter empty words and corresponding coordinates
							A_			    =		[idx for idx, word in enumerate(UpperCamelCase_ ) if not word.strip()]
							A_			    =		[word for idx, word in enumerate(UpperCamelCase_ ) if idx not in irrelevant_indices]
							A_			    =		[coord for idx, coord in enumerate(UpperCamelCase_ ) if idx not in irrelevant_indices]
							A_			    =		[coord for idx, coord in enumerate(UpperCamelCase_ ) if idx not in irrelevant_indices]
							A_			    =		[coord for idx, coord in enumerate(UpperCamelCase_ ) if idx not in irrelevant_indices]
							A_			    =		[coord for idx, coord in enumerate(UpperCamelCase_ ) if idx not in irrelevant_indices]

							# turn coordinates into (left, top, left+width, top+height) format
							A_			    =		[]
							for x, y, w, h in zip(UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_ ):
														A_			    =		[x, y, x + w, y + h]
														actual_boxes.append(UpperCamelCase_ )

							# finally, normalize the bounding boxes
							A_			    =		[]
							for box in actual_boxes:
														normalized_boxes.append(normalize_box(UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_ ) )

							assert len(UpperCamelCase_ ) == len(UpperCamelCase_ ), "Not as many words as there are bounding boxes"

							return words, normalized_boxes







class 					__lowerCAmelCase   (   lowercase						):

						"""simple docstring"""






						_UpperCAmelCase		:       Any   						=["pixel_values"]






						def __init__(  self							:      List[Any] ,  lowerCAmelCase							:      bool = True ,  lowerCAmelCase							:      Dict[str, int] = None ,  lowerCAmelCase							:      PILImageResampling = PILImageResampling.BILINEAR ,  lowerCAmelCase							:      bool = True ,  lowerCAmelCase							:      float = 1 / 2_55 ,  lowerCAmelCase							:      bool = True ,  lowerCAmelCase							:      Union[float, Iterable[float]] = None ,  lowerCAmelCase							:      Union[float, Iterable[float]] = None ,  lowerCAmelCase							:      bool = True ,  lowerCAmelCase							:      Optional[str] = None ,  lowerCAmelCase							:      Optional[str] = "" ,  **lowerCAmelCase							:      int ,  ):
													super().__init__(**lowerCAmelCase	)
													A_			    =		size if size is not None else {"height": 2_24, "width": 2_24}
													A_			    =		get_size_dict(lowerCAmelCase	)

													A_			    =		do_resize
													A_			    =		size
													A_			    =		resample
													A_			    =		do_rescale
													A_			    =		rescale_value
													A_			    =		do_normalize
													A_			    =		image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
													A_			    =		image_std if image_std is not None else IMAGENET_STANDARD_STD
													A_			    =		apply_ocr
													A_			    =		ocr_lang
													A_			    =		tesseract_config






						def 				_UpperCAmelCase       (  self							:      List[str] ,  lowerCAmelCase							:      np.ndarray ,  lowerCAmelCase							:      Dict[str, int] ,  lowerCAmelCase							:      PILImageResampling = PILImageResampling.BILINEAR ,  lowerCAmelCase							:      Optional[Union[str, ChannelDimension]] = None ,  **lowerCAmelCase							:      List[str] ,  ):
													A_			    =		get_size_dict(lowerCAmelCase	)
													if "height" not in size or "width" not in size:
																				raise ValueError(F"The size dictionary must contain the keys 'height' and 'width'. Got {size.keys()}"	)
													A_			    =		(size["height"], size["width"])
													return resize(lowerCAmelCase ,  size=lowerCAmelCase ,  resample=lowerCAmelCase ,  data_format=lowerCAmelCase ,  **lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      List[Any] ,  lowerCAmelCase							:      np.ndarray ,  lowerCAmelCase							:      Union[int, float] ,  lowerCAmelCase							:      Optional[Union[str, ChannelDimension]] = None ,  **lowerCAmelCase							:      Tuple ,  ):
													return rescale(lowerCAmelCase ,  scale=lowerCAmelCase ,  data_format=lowerCAmelCase ,  **lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      int ,  lowerCAmelCase							:      np.ndarray ,  lowerCAmelCase							:      Union[float, Iterable[float]] ,  lowerCAmelCase							:      Union[float, Iterable[float]] ,  lowerCAmelCase							:      Optional[Union[str, ChannelDimension]] = None ,  **lowerCAmelCase							:      int ,  ):
													return normalize(lowerCAmelCase ,  mean=lowerCAmelCase ,  std=lowerCAmelCase ,  data_format=lowerCAmelCase ,  **lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      Optional[Any] ,  lowerCAmelCase							:      ImageInput ,  lowerCAmelCase							:      bool = None ,  lowerCAmelCase							:      Dict[str, int] = None ,  lowerCAmelCase							:      Any=None ,  lowerCAmelCase							:      bool = None ,  lowerCAmelCase							:      float = None ,  lowerCAmelCase							:      bool = None ,  lowerCAmelCase							:      Union[float, Iterable[float]] = None ,  lowerCAmelCase							:      Union[float, Iterable[float]] = None ,  lowerCAmelCase							:      bool = None ,  lowerCAmelCase							:      Optional[str] = None ,  lowerCAmelCase							:      Optional[str] = None ,  lowerCAmelCase							:      Optional[Union[str, TensorType]] = None ,  lowerCAmelCase							:      ChannelDimension = ChannelDimension.FIRST ,  **lowerCAmelCase							:      int ,  ):
													A_			    =		do_resize if do_resize is not None else self.do_resize
													A_			    =		size if size is not None else self.size
													A_			    =		get_size_dict(lowerCAmelCase	)
													A_			    =		resample if resample is not None else self.resample
													A_			    =		do_rescale if do_rescale is not None else self.do_rescale
													A_			    =		rescale_factor if rescale_factor is not None else self.rescale_factor
													A_			    =		do_normalize if do_normalize is not None else self.do_normalize
													A_			    =		image_mean if image_mean is not None else self.image_mean
													A_			    =		image_std if image_std is not None else self.image_std
													A_			    =		apply_ocr if apply_ocr is not None else self.apply_ocr
													A_			    =		ocr_lang if ocr_lang is not None else self.ocr_lang
													A_			    =		tesseract_config if tesseract_config is not None else self.tesseract_config

													A_			    =		make_list_of_images(lowerCAmelCase	)

													if not valid_images(lowerCAmelCase	):
																				raise ValueError(
																				    "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
																				    "torch.Tensor, tf.Tensor or jax.ndarray."	)

													if do_resize and size is None:
																				raise ValueError("Size must be specified if do_resize is True."	)

													if do_rescale and rescale_factor is None:
																				raise ValueError("Rescale factor must be specified if do_rescale is True."	)

													if do_normalize and (image_mean is None or image_std is None):
																				raise ValueError("If do_normalize is True, image_mean and image_std must be specified."	)

													# All transformations expect numpy arrays.
													A_			    =		[to_numpy_array(lowerCAmelCase	) for image in images]

													# Tesseract OCR to get words + normalized bounding boxes
													if apply_ocr:
																				requires_backends(self ,  "pytesseract"	)
																				A_			    =		[]
																				A_			    =		[]
																				for image in images:
																											A_    ,					A_			    =		apply_tesseract(lowerCAmelCase ,  lowerCAmelCase ,  lowerCAmelCase	)
																											words_batch.append(lowerCAmelCase	)
																											boxes_batch.append(lowerCAmelCase	)

													if do_resize:
																				A_			    =		[self.resize(image=lowerCAmelCase ,  size=lowerCAmelCase ,  resample=lowerCAmelCase	) for image in images]

													if do_rescale:
																				A_			    =		[self.rescale(image=lowerCAmelCase ,  scale=lowerCAmelCase	) for image in images]

													if do_normalize:
																				A_			    =		[self.normalize(image=lowerCAmelCase ,  mean=lowerCAmelCase ,  std=lowerCAmelCase	) for image in images]

													A_			    =		[to_channel_dimension_format(lowerCAmelCase ,  lowerCAmelCase	) for image in images]

													A_			    =		BatchFeature(data={"pixel_values": images} ,  tensor_type=lowerCAmelCase	)

													if apply_ocr:
																				A_			    =		words_batch
																				A_			    =		boxes_batch
													return data





