usrbinenv python codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and finetuning any transformers model supported by automodelforsemanticsegmentation for semantic segmentation leveraging the trainer api logger logging getloggername will error if the minimal version of transformers is not installed remove at your own risks checkminversion4 36 0 dev0 requireversiondatasets2 0 0 to fix pip install r examplespytorchsemanticsegmentationrequirements txt def padifsmallerimg size fill0 size size size if isinstancesize int else size originalwidth originalheight img size padheight size1 originalheight if originalheight size1 else 0 padwidth size0 originalwidth if originalwidth size0 else 0 img functional padimg 0 0 padwidth padheight fillfill return img class compose def initself transforms self transforms transforms def callself image target for t in self transforms image target timage target return image target class identity def initself pass def callself image target return image target class resize def initself size self size size def callself image target image functional resizeimage self size target functional resizetarget self size interpolationtransforms interpolationmode nearest return image target class randomresize def initself minsize maxsizenone self minsize minsize if maxsize is none maxsize minsize self maxsize maxsize def callself image target size random randintself minsize self maxsize image functional resizeimage size target functional resizetarget size interpolationtransforms interpolationmode nearest return image target class randomcrop def initself size self size size if isinstancesize tuple else size size def callself image target image padifsmallerimage self size target padifsmallertarget self size fill255 cropparams transforms randomcrop getparamsimage self size image functional cropimage cropparams target functional croptarget cropparams return image target class randomhorizontalflip def initself flipprob self flipprob flipprob def callself image target if random random self flipprob image functional hflipimage target functional hfliptarget return image target class piltotensor def callself image target image functional piltotensorimage target torch astensornp arraytarget dtypetorch int64 return image target class convertimagedtype def initself dtype self dtype dtype def callself image target image functional convertimagedtypeimage self dtype return image target class normalize def initself mean std self mean mean self std std def callself image target image functional normalizeimage meanself mean stdself std return image target class reducelabels def callself image target if not isinstancetarget np ndarray target np arraytarget astypenp uint8 avoid using underflow conversion targettarget 0 255 target target 1 targettarget 254 255 target image fromarraytarget return image target dataclass class datatrainingarguments datasetname optionalstr field defaultsegmentssidewalksemantic metadata help name of a dataset from the hub could be your own possibly private dataset hosted on the hub datasetconfigname optionalstr field defaultnone metadatahelp the configuration name of the dataset to use via the datasets library trainvalsplit optionalfloat field default0 15 metadatahelp percent to split off of train for validation maxtrainsamples optionalint field defaultnone metadata help for debugging purposes or quicker training truncate the number of training examples to this value if set maxevalsamples optionalint field defaultnone metadata help for debugging purposes or quicker training truncate the number of evaluation examples to this value if set reducelabels optionalbool field defaultfalse metadatahelp whether or not to reduce all labels by 1 and replace background by 255 def postinitself if self datasetname is none and self traindir is none and self validationdir is none raise valueerror you must specify either a dataset name from the hub or a train andor validation directory dataclass class modelarguments modelnameorpath str field defaultnvidiamitb0 metadatahelp path to pretrained model or model identifier from huggingface comodels configname optionalstr field defaultnone metadatahelp pretrained config name or path if not the same as modelname cachedir optionalstr field defaultnone metadatahelp where do you want to store the pretrained models downloaded from s3 modelrevision str field defaultmain metadatahelp the specific model version to use can be a branch name tag name or commit id imageprocessorname str fielddefaultnone metadatahelp name or path of preprocessor config token str field defaultnone metadata help the token to use as http bearer ization for remote files if not specified will use the token generated when running huggingfacecli login stored in huggingface useauthtoken bool field defaultnone metadata help the useauthtoken argument is deprecated and will be removed in v4 34 please use token instead trustremotecode bool field defaultfalse metadata help whether or not to allow for custom models defined on the hub in their own modeling files this option should only be set to true for repositories you trust and in which you have read the code as it will execute code present on the hub on your local machine def main see all possible arguments in srctransformerstrainingargs py or by passing the help flag to this script we now keep distinct sets of args for a cleaner separation of concerns parser hfargumentparsermodelarguments datatrainingarguments trainingarguments if lensys argv 2 and sys argv1 endswith json if we pass only one argument to the script and it s the path to a json file let s parse it to get our arguments modelargs dataargs trainingargs parser parsejsonfilejsonfileos path abspathsys argv1 else modelargs dataargs trainingargs parser parseargsintodataclasses if modelargs useauthtoken is not none warnings warn the useauthtoken argument is deprecated and will be removed in v4 34 please use token instead futurewarning if modelargs token is not none raise valueerrortoken and useauthtoken are both specified please set only the argument token modelargs token modelargs useauthtoken sending telemetry tracking the example usage helps us better allocate resources to maintain them the information sent is the one passed as arguments along with your pythonpytorch versions sendexampletelemetryrunsemanticsegmentation modelargs dataargs setup logging logging basicconfig formatasctimes levelnames names messages datefmtmdy h m s handlerslogging streamhandlersys stdout if trainingargs shouldlog the default of trainingargs loglevel is passive so we set log level at info here to have that default transformers utils logging setverbosityinfo loglevel trainingargs getprocessloglevel logger setlevelloglevel transformers utils logging setverbosityloglevel transformers utils logging enabledefaulthandler transformers utils logging enableexplicitformat log on each process the small summary logger warning fprocess rank trainingargs localrank device trainingargs device ngpu trainingargs ngpu fdistributed training trainingargs parallelmode value distributed 16bits training trainingargs fp16 logger infoftrainingevaluation parameters trainingargs detecting last checkpoint lastcheckpoint none if os path isdirtrainingargs outputdir and trainingargs dotrain and not trainingargs overwriteoutputdir lastcheckpoint getlastcheckpointtrainingargs outputdir if lastcheckpoint is none and lenos listdirtrainingargs outputdir 0 raise valueerror foutput directory trainingargs outputdir already exists and is not empty use overwriteoutputdir to overcome elif lastcheckpoint is not none and trainingargs resumefromcheckpoint is none logger info fcheckpoint detected resuming training at lastcheckpoint to avoid this behavior change the outputdir or add overwriteoutputdir to train from scratch load dataset in distributed training the loaddataset function guarantees that only one local process can concurrently download the dataset todo support datasets from local folders dataset loaddatasetdataargs datasetname cachedirmodelargs cachedir rename column names to standardized names only image and label need to be present if pixelvalues in datasettrain columnnames dataset dataset renamecolumnspixelvalues image if annotation in datasettrain columnnames dataset dataset renamecolumnsannotation label if we don t have a validation split split off a percentage of train as validation dataargs trainvalsplit none if validation in dataset keys else dataargs trainvalsplit if isinstancedataargs trainvalsplit float and dataargs trainvalsplit 0 0 split datasettrain traintestsplitdataargs trainvalsplit datasettrain splittrain datasetvalidation splittest prepare label mappings we ll include these in the model s config to get human readable labels in the inference api if dataargs datasetname sceneparse150 repoid huggingfacelabelfiles filename ade20kid2label json else repoid dataargs datasetname filename id2label json id2label json loadopenhfhubdownloadrepoid filename repotypedataset r id2label intk v for k v in id2label items label2id v strk for k v in id2label items load the mean iou metric from the datasets package metric evaluate loadmeaniou define our computemetrics function it takes an evalprediction object a namedtuple with a predictions and labelids field and has to return a dictionary string to float torch nograd def computemetricsevalpred logits labels evalpred logitstensor torch fromnumpylogits scale the logits to the size of the label logitstensor nn functional interpolate logitstensor sizelabels shape2 modebilinear aligncornersfalse argmaxdim1 predlabels logitstensor detach cpu numpy metrics metric compute predictionspredlabels referenceslabels numlabelslenid2label ignoreindex0 reducelabelsimageprocessor doreducelabels add per category metrics as individual keyvalue pairs percategoryaccuracy metrics poppercategoryaccuracy tolist percategoryiou metrics poppercategoryiou tolist metrics updatefaccuracyid2labeli v for i v in enumeratepercategoryaccuracy metrics updatefiouid2labeli v for i v in enumeratepercategoryiou return metrics config autoconfig frompretrained modelargs configname or modelargs modelnameorpath label2idlabel2id id2labelid2label cachedirmodelargs cachedir revisionmodelargs modelrevision tokenmodelargs token trustremotecodemodelargs trustremotecode model automodelforsemanticsegmentation frompretrained modelargs modelnameorpath fromtfbool ckpt in modelargs modelnameorpath configconfig cachedirmodelargs cachedir revisionmodelargs modelrevision tokenmodelargs token trustremotecodemodelargs trustremotecode imageprocessor autoimageprocessor frompretrained modelargs imageprocessorname or modelargs modelnameorpath cachedirmodelargs cachedir revisionmodelargs modelrevision tokenmodelargs token trustremotecodemodelargs trustremotecode define torchvision transforms to be applied to each image target not that straightforward in torchvision https github compytorchvisionissues9 currently based on official torchvision references https github compytorchvisionblobmainreferencessegmentationtransforms py if shortestedge in imageprocessor size we instead set the target size as shortestedge shortestedge to here to ensure all images are batchable size imageprocessor sizeshortestedge imageprocessor sizeshortestedge else size imageprocessor sizeheight imageprocessor sizewidth traintransforms compose reducelabels if dataargs reducelabels else identity randomcropsizesize randomhorizontalflipflipprob0 5 piltotensor convertimagedtypetorch float normalizemeanimageprocessor imagemean stdimageprocessor imagestd define torchvision transform to be applied to each image jitter colorjitterbrightness0 25 contrast0 25 saturation0 25 hue0 1 valtransforms compose reducelabels if dataargs reducelabels else identity resizesizesize piltotensor convertimagedtypetorch float normalizemeanimageprocessor imagemean stdimageprocessor imagestd def preprocesstrainexamplebatch pixelvalues labels for image target in zipexamplebatchimage examplebatchlabel image target traintransformsimage convertrgb target pixelvalues appendimage labels appendtarget encoding encodingpixelvalues torch stackpixelvalues encodinglabels torch stacklabels return encoding def preprocessvalexamplebatch pixelvalues labels for image target in zipexamplebatchimage examplebatchlabel image target valtransformsimage convertrgb target pixelvalues appendimage labels appendtarget encoding encodingpixelvalues torch stackpixelvalues encodinglabels torch stacklabels return encoding if trainingargs dotrain if train not in dataset raise valueerrordotrain requires a train dataset if dataargs maxtrainsamples is not none datasettrain datasettrain shuffleseedtrainingargs seed selectrangedataargs maxtrainsamples set the training transforms datasettrain settransformpreprocesstrain if trainingargs doeval if validation not in dataset raise valueerrordoeval requires a validation dataset if dataargs maxevalsamples is not none datasetvalidation datasetvalidation shuffleseedtrainingargs seed selectrangedataargs maxevalsamples set the validation transforms datasetvalidation settransformpreprocessval initalize our trainer trainer trainer modelmodel argstrainingargs traindatasetdatasettrain if trainingargs dotrain else none evaldatasetdatasetvalidation if trainingargs doeval else none computemetricscomputemetrics tokenizerimageprocessor datacollatordefaultdatacollator training if trainingargs dotrain checkpoint none if trainingargs resumefromcheckpoint is not none checkpoint trainingargs resumefromcheckpoint elif lastcheckpoint is not none checkpoint lastcheckpoint trainresult trainer trainresumefromcheckpointcheckpoint trainer savemodel trainer logmetricstrain trainresult metrics trainer savemetricstrain trainresult metrics trainer savestate evaluation if trainingargs doeval metrics trainer evaluate trainer logmetricseval metrics trainer savemetricseval metrics write model card and optionally push to hub kwargs finetunedfrom modelargs modelnameorpath dataset dataargs datasetname tags imagesegmentation vision if trainingargs pushtohub trainer pushtohubkwargs else trainer createmodelcardkwargs if name main main usr bin env python coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and finetuning any transformers model supported by automodelforsemanticsegmentation for semantic segmentation leveraging the trainer api will error if the minimal version of transformers is not installed remove at your own risks avoid using underflow conversion arguments pertaining to what data we are going to input our model for training and eval using hfargumentparser we can turn this class into argparse arguments to be able to specify them on the command line arguments pertaining to which model config tokenizer we are going to fine tune from see all possible arguments in src transformers training_args py or by passing the help flag to this script we now keep distinct sets of args for a cleaner separation of concerns if we pass only one argument to the script and it s the path to a json file let s parse it to get our arguments sending telemetry tracking the example usage helps us better allocate resources to maintain them the information sent is the one passed as arguments along with your python pytorch versions setup logging the default of training_args log_level is passive so we set log level at info here to have that default log on each process the small summary detecting last checkpoint load dataset in distributed training the load_dataset function guarantees that only one local process can concurrently download the dataset todo support datasets from local folders rename column names to standardized names only image and label need to be present if we don t have a validation split split off a percentage of train as validation prepare label mappings we ll include these in the model s config to get human readable labels in the inference api load the mean iou metric from the datasets package define our compute_metrics function it takes an evalprediction object a namedtuple with a predictions and label_ids field and has to return a dictionary string to float scale the logits to the size of the label add per category metrics as individual key value pairs define torchvision transforms to be applied to each image target not that straightforward in torchvision https github com pytorch vision issues 9 currently based on official torchvision references https github com pytorch vision blob main references segmentation transforms py we instead set the target size as shortest_edge shortest_edge to here to ensure all images are batchable define torchvision transform to be applied to each image jitter colorjitter brightness 0 25 contrast 0 25 saturation 0 25 hue 0 1 set the training transforms set the validation transforms initalize our trainer training evaluation write model card and optionally push to hub
import json
import logging
import os
import random
import sys
import warnings
from dataclasses import dataclass, field
from typing import Optional

import evaluate
import numpy as np
import torch
from datasets import load_dataset
from huggingface_hub import hf_hub_download
from PIL import Image
from torch import nn
from torchvision import transforms
from torchvision.transforms import functional

import transformers
from transformers import (
    AutoConfig,
    AutoImageProcessor,
    AutoModelForSemanticSegmentation,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    default_data_collator,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version




logger = logging.getLogger(__name__)


check_min_version("4.36.0.dev0")

require_version("datasets>=2.0.0", "To fix: pip install -r examples/pytorch/semantic-segmentation/requirements.txt")


def pad_if_smaller(img, size, fill=0):
    size = (size, size) if isinstance(size, int) else size
    original_width, original_height = img.size
    pad_height = size[1] - original_height if original_height < size[1] else 0
    pad_width = size[0] - original_width if original_width < size[0] else 0
    img = functional.pad(img, (0, 0, pad_width, pad_height), fill=fill)
    return img


class Compose:
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image, target):
        for t in self.transforms:
            image, target = t(image, target)
        return image, target


class Identity:
    def __init__(self):
        pass

    def __call__(self, image, target):
        return image, target


class Resize:
    def __init__(self, size):
        self.size = size

    def __call__(self, image, target):
        image = functional.resize(image, self.size)
        target = functional.resize(target, self.size, interpolation=transforms.InterpolationMode.NEAREST)
        return image, target


class RandomResize:
    def __init__(self, min_size, max_size=None):
        self.min_size = min_size
        if max_size is None:
            max_size = min_size
        self.max_size = max_size

    def __call__(self, image, target):
        size = random.randint(self.min_size, self.max_size)
        image = functional.resize(image, size)
        target = functional.resize(target, size, interpolation=transforms.InterpolationMode.NEAREST)
        return image, target


class RandomCrop:
    def __init__(self, size):
        self.size = size if isinstance(size, tuple) else (size, size)

    def __call__(self, image, target):
        image = pad_if_smaller(image, self.size)
        target = pad_if_smaller(target, self.size, fill=255)
        crop_params = transforms.RandomCrop.get_params(image, self.size)
        image = functional.crop(image, *crop_params)
        target = functional.crop(target, *crop_params)
        return image, target


class RandomHorizontalFlip:
    def __init__(self, flip_prob):
        self.flip_prob = flip_prob

    def __call__(self, image, target):
        if random.random() < self.flip_prob:
            image = functional.hflip(image)
            target = functional.hflip(target)
        return image, target


class PILToTensor:
    def __call__(self, image, target):
        image = functional.pil_to_tensor(image)
        target = torch.as_tensor(np.array(target), dtype=torch.int64)
        return image, target


class ConvertImageDtype:
    def __init__(self, dtype):
        self.dtype = dtype

    def __call__(self, image, target):
        image = functional.convert_image_dtype(image, self.dtype)
        return image, target


class Normalize:
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, image, target):
        image = functional.normalize(image, mean=self.mean, std=self.std)
        return image, target


class ReduceLabels:
    def __call__(self, image, target):
        if not isinstance(target, np.ndarray):
            target = np.array(target).astype(np.uint8)
        
        target[target == 0] = 255
        target = target - 1
        target[target == 254] = 255

        target = Image.fromarray(target)
        return image, target


@dataclass
class DataTrainingArguments:
    

    dataset_name: Optional[str] = field(
        default="segments/sidewalk-semantic",
        metadata={
            "help": "Name of a dataset from the hub (could be your own, possibly private dataset hosted on the hub)."
        },
    )
    dataset_config_name: Optional[str] = field(
        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
    )
    train_val_split: Optional[float] = field(
        default=0.15, metadata={"help": "Percent to split off of train for validation."}
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of training examples to this "
                "value if set."
            )
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
                "value if set."
            )
        },
    )
    reduce_labels: Optional[bool] = field(
        default=False,
        metadata={"help": "Whether or not to reduce all labels by 1 and replace background by 255."},
    )

    def __post_init__(self):
        if self.dataset_name is None and (self.train_dir is None and self.validation_dir is None):
            raise ValueError(
                "You must specify either a dataset name from the hub or a train and/or validation directory."
            )


@dataclass
class ModelArguments:
    

    model_name_or_path: str = field(
        default="nvidia/mit-b0",
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"},
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
    )
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
    image_processor_name: str = field(default=None, metadata={"help": "Name or path of preprocessor config."})
    token: str = field(
        default=None,
        metadata={
            "help": (
                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
            )
        },
    )
    use_auth_token: bool = field(
        default=None,
        metadata={
            "help": "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token` instead."
        },
    )
    trust_remote_code: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether or not to allow for custom models defined on the Hub in their own modeling files. This option"
                "should only be set to `True` for repositories you trust and in which you have read the code, as it will "
                "execute code present on the Hub on your local machine."
            )
        },
    )


def main():
    
    
    

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        
        
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    if model_args.use_auth_token is not None:
        warnings.warn(
            "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token` instead.",
            FutureWarning,
        )
        if model_args.token is not None:
            raise ValueError("`token` and `use_auth_token` are both specified. Please set only the argument `token`.")
        model_args.token = model_args.use_auth_token

    
    
    send_example_telemetry("run_semantic_segmentation", model_args, data_args)

    
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )

    if training_args.should_log:
        
        transformers.utils.logging.set_verbosity_info()

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
        + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    
    
    
    
    dataset = load_dataset(data_args.dataset_name, cache_dir=model_args.cache_dir)

    
    if "pixel_values" in dataset["train"].column_names:
        dataset = dataset.rename_columns({"pixel_values": "image"})
    if "annotation" in dataset["train"].column_names:
        dataset = dataset.rename_columns({"annotation": "label"})

    
    data_args.train_val_split = None if "validation" in dataset.keys() else data_args.train_val_split
    if isinstance(data_args.train_val_split, float) and data_args.train_val_split > 0.0:
        split = dataset["train"].train_test_split(data_args.train_val_split)
        dataset["train"] = split["train"]
        dataset["validation"] = split["test"]

    
    
    if data_args.dataset_name == "scene_parse_150":
        repo_id = "huggingface/label-files"
        filename = "ade20k-id2label.json"
    else:
        repo_id = data_args.dataset_name
        filename = "id2label.json"
    id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
    id2label = {int(k): v for k, v in id2label.items()}
    label2id = {v: str(k) for k, v in id2label.items()}

    
    metric = evaluate.load("mean_iou")

    
    
    @torch.no_grad()
    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        logits_tensor = torch.from_numpy(logits)
        
        logits_tensor = nn.functional.interpolate(
            logits_tensor,
            size=labels.shape[-2:],
            mode="bilinear",
            align_corners=False,
        ).argmax(dim=1)

        pred_labels = logits_tensor.detach().cpu().numpy()
        metrics = metric.compute(
            predictions=pred_labels,
            references=labels,
            num_labels=len(id2label),
            ignore_index=0,
            reduce_labels=image_processor.do_reduce_labels,
        )
        
        per_category_accuracy = metrics.pop("per_category_accuracy").tolist()
        per_category_iou = metrics.pop("per_category_iou").tolist()

        metrics.update({f"accuracy_{id2label[i]}": v for i, v in enumerate(per_category_accuracy)})
        metrics.update({f"iou_{id2label[i]}": v for i, v in enumerate(per_category_iou)})

        return metrics

    config = AutoConfig.from_pretrained(
        model_args.config_name or model_args.model_name_or_path,
        label2id=label2id,
        id2label=id2label,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=model_args.token,
        trust_remote_code=model_args.trust_remote_code,
    )
    model = AutoModelForSemanticSegmentation.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=model_args.token,
        trust_remote_code=model_args.trust_remote_code,
    )
    image_processor = AutoImageProcessor.from_pretrained(
        model_args.image_processor_name or model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=model_args.token,
        trust_remote_code=model_args.trust_remote_code,
    )

    
    
    
    if "shortest_edge" in image_processor.size:
        
        size = (image_processor.size["shortest_edge"], image_processor.size["shortest_edge"])
    else:
        size = (image_processor.size["height"], image_processor.size["width"])
    train_transforms = Compose(
        [
            ReduceLabels() if data_args.reduce_labels else Identity(),
            RandomCrop(size=size),
            RandomHorizontalFlip(flip_prob=0.5),
            PILToTensor(),
            ConvertImageDtype(torch.float),
            Normalize(mean=image_processor.image_mean, std=image_processor.image_std),
        ]
    )
    
    
    val_transforms = Compose(
        [
            ReduceLabels() if data_args.reduce_labels else Identity(),
            Resize(size=size),
            PILToTensor(),
            ConvertImageDtype(torch.float),
            Normalize(mean=image_processor.image_mean, std=image_processor.image_std),
        ]
    )

    def preprocess_train(example_batch):
        pixel_values = []
        labels = []
        for image, target in zip(example_batch["image"], example_batch["label"]):
            image, target = train_transforms(image.convert("RGB"), target)
            pixel_values.append(image)
            labels.append(target)

        encoding = {}
        encoding["pixel_values"] = torch.stack(pixel_values)
        encoding["labels"] = torch.stack(labels)

        return encoding

    def preprocess_val(example_batch):
        pixel_values = []
        labels = []
        for image, target in zip(example_batch["image"], example_batch["label"]):
            image, target = val_transforms(image.convert("RGB"), target)
            pixel_values.append(image)
            labels.append(target)

        encoding = {}
        encoding["pixel_values"] = torch.stack(pixel_values)
        encoding["labels"] = torch.stack(labels)

        return encoding

    if training_args.do_train:
        if "train" not in dataset:
            raise ValueError("--do_train requires a train dataset")
        if data_args.max_train_samples is not None:
            dataset["train"] = (
                dataset["train"].shuffle(seed=training_args.seed).select(range(data_args.max_train_samples))
            )
        
        dataset["train"].set_transform(preprocess_train)

    if training_args.do_eval:
        if "validation" not in dataset:
            raise ValueError("--do_eval requires a validation dataset")
        if data_args.max_eval_samples is not None:
            dataset["validation"] = (
                dataset["validation"].shuffle(seed=training_args.seed).select(range(data_args.max_eval_samples))
            )
        
        dataset["validation"].set_transform(preprocess_val)

    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"] if training_args.do_train else None,
        eval_dataset=dataset["validation"] if training_args.do_eval else None,
        compute_metrics=compute_metrics,
        tokenizer=image_processor,
        data_collator=default_data_collator,
    )

    
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model()
        trainer.log_metrics("train", train_result.metrics)
        trainer.save_metrics("train", train_result.metrics)
        trainer.save_state()

    
    if training_args.do_eval:
        metrics = trainer.evaluate()
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)

    
    kwargs = {
        "finetuned_from": model_args.model_name_or_path,
        "dataset": data_args.dataset_name,
        "tags": ["image-segmentation", "vision"],
    }
    if training_args.push_to_hub:
        trainer.push_to_hub(**kwargs)
    else:
        trainer.create_model_card(**kwargs)


if __name__ == "__main__":
    main()