coreml-community/ControlNet-v1-1-Annotators-cpu

1import os.path as osp2import pickle3import shutil4import tempfile5 6import annotator.mmpkg.mmcv as mmcv7import numpy as np8import torch9import torch.distributed as dist10from annotator.mmpkg.mmcv.image import tensor2imgs11from annotator.mmpkg.mmcv.runner import get_dist_info12 13 14def np2tmp(array, temp_file_name=None):15    """Save ndarray to local numpy file.16 17    Args:18        array (ndarray): Ndarray to save.19        temp_file_name (str): Numpy file name. If 'temp_file_name=None', this20            function will generate a file name with tempfile.NamedTemporaryFile21            to save ndarray. Default: None.22 23    Returns:24        str: The numpy file name.25    """26 27    if temp_file_name is None:28        temp_file_name = tempfile.NamedTemporaryFile(29            suffix='.npy', delete=False).name30    np.save(temp_file_name, array)31    return temp_file_name32 33 34def single_gpu_test(model,35                    data_loader,36                    show=False,37                    out_dir=None,38                    efficient_test=False,39                    opacity=0.5):40    """Test with single GPU.41 42    Args:43        model (nn.Module): Model to be tested.44        data_loader (utils.data.Dataloader): Pytorch data loader.45        show (bool): Whether show results during inference. Default: False.46        out_dir (str, optional): If specified, the results will be dumped into47            the directory to save output results.48        efficient_test (bool): Whether save the results as local numpy files to49            save CPU memory during evaluation. Default: False.50        opacity(float): Opacity of painted segmentation map.51            Default 0.5.52            Must be in (0, 1] range.53    Returns:54        list: The prediction results.55    """56 57    model.eval()58    results = []59    dataset = data_loader.dataset60    prog_bar = mmcv.ProgressBar(len(dataset))61    for i, data in enumerate(data_loader):62        with torch.no_grad():63            result = model(return_loss=False, **data)64 65        if show or out_dir:66            img_tensor = data['img'][0]67            img_metas = data['img_metas'][0].data[0]68            imgs = tensor2imgs(img_tensor, **img_metas[0]['img_norm_cfg'])69            assert len(imgs) == len(img_metas)70 71            for img, img_meta in zip(imgs, img_metas):72                h, w, _ = img_meta['img_shape']73                img_show = img[:h, :w, :]74 75                ori_h, ori_w = img_meta['ori_shape'][:-1]76                img_show = mmcv.imresize(img_show, (ori_w, ori_h))77 78                if out_dir:79                    out_file = osp.join(out_dir, img_meta['ori_filename'])80                else:81                    out_file = None82 83                model.module.show_result(84                    img_show,85                    result,86                    palette=dataset.PALETTE,87                    show=show,88                    out_file=out_file,89                    opacity=opacity)90 91        if isinstance(result, list):92            if efficient_test:93                result = [np2tmp(_) for _ in result]94            results.extend(result)95        else:96            if efficient_test:97                result = np2tmp(result)98            results.append(result)99 100        batch_size = len(result)101        for _ in range(batch_size):102            prog_bar.update()103    return results104 105 106def multi_gpu_test(model,107                   data_loader,108                   tmpdir=None,109                   gpu_collect=False,110                   efficient_test=False):111    """Test model with multiple gpus.112 113    This method tests model with multiple gpus and collects the results114    under two different modes: gpu and cpu modes. By setting 'gpu_collect=True'115    it encodes results to gpu tensors and use gpu communication for results116    collection. On cpu mode it saves the results on different gpus to 'tmpdir'117    and collects them by the rank 0 worker.118 119    Args:120        model (nn.Module): Model to be tested.121        data_loader (utils.data.Dataloader): Pytorch data loader.122        tmpdir (str): Path of directory to save the temporary results from123            different gpus under cpu mode.124        gpu_collect (bool): Option to use either gpu or cpu to collect results.125        efficient_test (bool): Whether save the results as local numpy files to126            save CPU memory during evaluation. Default: False.127 128    Returns:129        list: The prediction results.130    """131 132    model.eval()133    results = []134    dataset = data_loader.dataset135    rank, world_size = get_dist_info()136    if rank == 0:137        prog_bar = mmcv.ProgressBar(len(dataset))138    for i, data in enumerate(data_loader):139        with torch.no_grad():140            result = model(return_loss=False, rescale=True, **data)141 142        if isinstance(result, list):143            if efficient_test:144                result = [np2tmp(_) for _ in result]145            results.extend(result)146        else:147            if efficient_test:148                result = np2tmp(result)149            results.append(result)150 151        if rank == 0:152            batch_size = data['img'][0].size(0)153            for _ in range(batch_size * world_size):154                prog_bar.update()155 156    # collect results from all ranks157    if gpu_collect:158        results = collect_results_gpu(results, len(dataset))159    else:160        results = collect_results_cpu(results, len(dataset), tmpdir)161    return results162 163 164def collect_results_cpu(result_part, size, tmpdir=None):165    """Collect results with CPU."""166    rank, world_size = get_dist_info()167    # create a tmp dir if it is not specified168    if tmpdir is None:169        MAX_LEN = 512170        # 32 is whitespace171        dir_tensor = torch.full((MAX_LEN, ),172                                32,173                                dtype=torch.uint8,174                                device='cuda')175        if rank == 0:176            tmpdir = tempfile.mkdtemp()177            tmpdir = torch.tensor(178                bytearray(tmpdir.encode()), dtype=torch.uint8, device='cuda')179            dir_tensor[:len(tmpdir)] = tmpdir180        dist.broadcast(dir_tensor, 0)181        tmpdir = dir_tensor.cpu().numpy().tobytes().decode().rstrip()182    else:183        mmcv.mkdir_or_exist(tmpdir)184    # dump the part result to the dir185    mmcv.dump(result_part, osp.join(tmpdir, 'part_{}.pkl'.format(rank)))186    dist.barrier()187    # collect all parts188    if rank != 0:189        return None190    else:191        # load results of all parts from tmp dir192        part_list = []193        for i in range(world_size):194            part_file = osp.join(tmpdir, 'part_{}.pkl'.format(i))195            part_list.append(mmcv.load(part_file))196        # sort the results197        ordered_results = []198        for res in zip(*part_list):199            ordered_results.extend(list(res))200        # the dataloader may pad some samples201        ordered_results = ordered_results[:size]202        # remove tmp dir203        shutil.rmtree(tmpdir)204        return ordered_results205 206 207def collect_results_gpu(result_part, size):208    """Collect results with GPU."""209    rank, world_size = get_dist_info()210    # dump result part to tensor with pickle211    part_tensor = torch.tensor(212        bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device='cuda')213    # gather all result part tensor shape214    shape_tensor = torch.tensor(part_tensor.shape, device='cuda')215    shape_list = [shape_tensor.clone() for _ in range(world_size)]216    dist.all_gather(shape_list, shape_tensor)217    # padding result part tensor to max length218    shape_max = torch.tensor(shape_list).max()219    part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda')220    part_send[:shape_tensor[0]] = part_tensor221    part_recv_list = [222        part_tensor.new_zeros(shape_max) for _ in range(world_size)223    ]224    # gather all result part225    dist.all_gather(part_recv_list, part_send)226 227    if rank == 0:228        part_list = []229        for recv, shape in zip(part_recv_list, shape_list):230            part_list.append(231                pickle.loads(recv[:shape[0]].cpu().numpy().tobytes()))232        # sort the results233        ordered_results = []234        for res in zip(*part_list):235            ordered_results.extend(list(res))236        # the dataloader may pad some samples237        ordered_results = ordered_results[:size]238        return ordered_results239