3 Nb @sndZddlZddlZddlZddlZyddlmZWn ek rTddl mZYnXd ddZ d dd Z dS) zJModified from https://github.com/CSAILVision/semantic-segmentation-pytorchN) urlretrieve ./pretrainedcCsjtjj|stj||jdd}tjj||}tjj|s\tjjdj ||t ||t j ||dS)N/zDownloading: "{}" to {} ) map_location) ospathexistsmakedirssplitjoinsysstderrwriteformatrtorchload)url model_dirrfilename cached_filer4/home/ec2-user/SageMaker/lama/models/ade20k/utils.pyload_urls    rRGBcCs|jd}tj|jd|jddftjd}x^tj|D]P}|dkrFq8|||kddddtjftj|||jd|jddf7}q8W|dkr|dddddddfS|SdS)Nintrr)dtypeZBGRr)astypenpzerosshapeuint8uniquenewaxistile)ZlabelmapcolorsmodeZ labelmap_rgblabelrrr color_encodes    r*)rN)r) __doc__rrnumpyr rurllibr ImportErrorurllib.requestrr*rrrrs