Esempio n. 1
0
    def __init__(self):
        """
        extract semantic map from pretrained model, for knowing where to ignore in the image,
        since we only have depth info on mountains.
        """
        self.names = {}
        with open('semseg/object150_info.csv') as f:
            reader = csv.reader(f)
            next(reader)
            for row in reader:
                self.names[int(row[0])] = row[5].split(";")[0]

        # Network Builders
        self.net_encoder = ModelBuilder.build_encoder(
            arch='resnet50dilated',
            fc_dim=2048,
            weights=
            'semseg/ckpt/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth'
        )
        self.net_decoder = ModelBuilder.build_decoder(
            arch='ppm_deepsup',
            fc_dim=2048,
            num_class=150,
            weights=
            'semseg/ckpt/ade20k-resnet50dilated-ppm_deepsup/decoder_epoch_20.pth',
            use_softmax=True)

        self.crit = torch.nn.NLLLoss(ignore_index=-1)
        self.segmentation_module = SegmentationModule(self.net_encoder,
                                                      self.net_decoder,
                                                      self.crit)
        self.segmentation_module.eval()
        self.segmentation_module.to(device=defs.get_dev())
Esempio n. 2
0
def initialize():
    colors = scipy.io.loadmat('data/color150.mat')['colors']
    names = {}
    with open('data/object150_info.csv') as f:
        reader = csv.reader(f)
        next(reader)
        for row in reader:
            names[int(row[0])] = row[5].split(";")[0]

    # Network Builders
    net_encoder = ModelBuilder.build_encoder(
        arch='resnet50dilated',
        fc_dim=2048,
        weights='ckpt/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth')
    net_decoder = ModelBuilder.build_decoder(
        arch='ppm_deepsup',
        fc_dim=2048,
        num_class=150,
        weights='ckpt/ade20k-resnet50dilated-ppm_deepsup/decoder_epoch_20.pth',
        use_softmax=True)

    crit = torch.nn.NLLLoss(ignore_index=-1)
    segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
    segmentation_module = segmentation_module.eval()

    x = torch.rand(3,672,504)
    output_size = x.shape

    # segmentation_module.cuda()
    with torch.no_grad():
        with torch.jit.optimized_execution(True, {'target_device': 'eia:0'}):
            segmentation_module = torch.jit.trace(segmentation_module, x[None])
#         torch.jit.save(segmentation_module, 'traced.pt')
    return segmentation_module
    def _load_mistcsail_model(encoder, decoder, encoder_path=None, decoder_path=None):
        from mit_semseg.models import ModelBuilder, SegmentationModule
        from mit_semseg.config import cfg
        if encoder_path is None:
            encoder_path = str(Path(__file__).parents[0].joinpath('weights', f'encoder_{encoder}.pth'))
        if decoder_path is None:
            decoder_path = str(Path(__file__).parents[0].joinpath('weights', f'decoder_{decoder}.pth'))

        config_file = str(Path(__file__).parents[0].joinpath('config', f'ade20k-{encoder}-{decoder}.yaml'))

        cfg.merge_from_file(config_file)

        net_encoder = ModelBuilder.build_encoder(
            arch=encoder,
            fc_dim=cfg.MODEL.fc_dim,
            weights=encoder_path)
        net_decoder = ModelBuilder.build_decoder(
            arch=decoder,
            fc_dim=cfg.MODEL.fc_dim,
            num_class=cfg.DATASET.num_class,
            weights=decoder_path,
            use_softmax=True)

        crit = torch.nn.NLLLoss(ignore_index=-1)
        model = SegmentationModule(net_encoder, net_decoder, crit)
        return model
def worker(cfg, gpu_id, start_idx, end_idx, result_queue):
    torch.cuda.set_device(gpu_id)

    # Dataset and Loader
    dataset_val = ValDataset(cfg.DATASET.root_dataset,
                             cfg.DATASET.list_val,
                             cfg.DATASET,
                             start_idx=start_idx,
                             end_idx=end_idx)
    loader_val = torch.utils.data.DataLoader(dataset_val,
                                             batch_size=cfg.VAL.batch_size,
                                             shuffle=False,
                                             collate_fn=user_scattered_collate,
                                             num_workers=2)

    # Network Builders
    net_encoder = ModelBuilder.build_encoder(
        arch=cfg.MODEL.arch_encoder.lower(),
        fc_dim=cfg.MODEL.fc_dim,
        weights=cfg.MODEL.weights_encoder)
    net_decoder = ModelBuilder.build_decoder(
        arch=cfg.MODEL.arch_decoder.lower(),
        fc_dim=cfg.MODEL.fc_dim,
        num_class=cfg.DATASET.num_class,
        weights=cfg.MODEL.weights_decoder,
        use_softmax=True)

    crit = nn.NLLLoss(ignore_index=-1)

    segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)

    segmentation_module.cuda()

    # Main loop
    evaluate(segmentation_module, loader_val, cfg, gpu_id, result_queue)
Esempio n. 5
0
def main(cfg, gpu):
    torch.cuda.set_device(gpu)

    # Network Builders
    net_encoder = ModelBuilder.build_encoder(arch=cfg.MODEL.arch_encoder,
                                             fc_dim=cfg.MODEL.fc_dim,
                                             weights=cfg.MODEL.weights_encoder)
    net_decoder = ModelBuilder.build_decoder(arch=cfg.MODEL.arch_decoder,
                                             fc_dim=cfg.MODEL.fc_dim,
                                             num_class=cfg.DATASET.num_class,
                                             weights=cfg.MODEL.weights_decoder,
                                             use_softmax=True)

    crit = nn.NLLLoss(ignore_index=-1)

    segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)

    # Dataset and Loader
    dataset_test = TestDataset(cfg.list_test, cfg.DATASET)
    loader_test = torch.utils.data.DataLoader(
        dataset_test,
        batch_size=cfg.TEST.batch_size,
        shuffle=False,
        collate_fn=user_scattered_collate,
        num_workers=5,
        drop_last=True)

    segmentation_module.cuda()

    # Main loop
    test(segmentation_module, loader_test, gpu)

    print('Inference done!')
Esempio n. 6
0
def setup(gpu):
    global segmentation_module

    # Network Builders
    net_encoder = ModelBuilder.build_encoder(
        arch='resnet50dilated',
        fc_dim=2048,
        weights=downloads.download_data_file(
            'http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth',
            'semantic-segmentation-pytorch/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth'
        ))

    net_decoder = ModelBuilder.build_decoder(
        arch='ppm_deepsup',
        fc_dim=2048,
        num_class=150,
        weights=downloads.download_data_file(
            'http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-resnet50dilated-ppm_deepsup/decoder_epoch_20.pth',
            'semantic-segmentation-pytorch/ade20k-resnet50dilated-ppm_deepsup/decoder_epoch_20.pth'
        ),
        use_softmax=True)

    crit = torch.nn.NLLLoss(ignore_index=-1)
    segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
    segmentation_module.eval()
    segmentation_module.cuda()
Esempio n. 7
0
    def _load_model(self, cfg_file, gpu):
        if gpu is not None:
            torch.cuda.set_device(0)
        basepath = rospkg.RosPack().get_path('object_segmentation')
        cfg.merge_from_file(basepath + "/" + cfg_file)

        logger = setup_logger(distributed_rank=0)
        logger.info(f"Loaded configuration file {cfg_file}")
        logger.info("Running with config:\n{}".format(cfg))

        cfg.MODEL.arch_encoder = cfg.MODEL.arch_encoder.lower()
        cfg.MODEL.arch_decoder = cfg.MODEL.arch_decoder.lower()

        # absolute paths of model weights
        cfg.MODEL.weights_encoder = (
            Path(basepath) / cfg.DIR /
            ('encoder_' + cfg.TEST.checkpoint)).as_posix()
        cfg.MODEL.weights_decoder = (
            Path(basepath) / cfg.DIR /
            ('decoder_' + cfg.TEST.checkpoint)).as_posix()

        if not os.path.exists(cfg.MODEL.weights_encoder) or not os.path.exists(
                cfg.MODEL.weights_decoder):
            download.ycb(Path(basepath) / 'ckpt')

        assert os.path.exists(
            cfg.MODEL.weights_encoder
        ), f"checkpoint {cfg.MODEL.weights_encoder} does not exitst!"
        assert os.path.exists(
            cfg.MODEL.weights_decoder
        ), f"checkpoint {cfg.MODEL.weights_decoder} does not exitst!"

        # Network Builders
        net_encoder = ModelBuilder.build_encoder(
            arch=cfg.MODEL.arch_encoder,
            fc_dim=cfg.MODEL.fc_dim,
            weights=cfg.MODEL.weights_encoder)
        net_decoder = ModelBuilder.build_decoder(
            arch=cfg.MODEL.arch_decoder,
            fc_dim=cfg.MODEL.fc_dim,
            num_class=cfg.DATASET.num_class,
            weights=cfg.MODEL.weights_decoder,
            use_softmax=True)

        crit = nn.NLLLoss(ignore_index=-1)

        segmentation_module = SegmentationModule(net_encoder, net_decoder,
                                                 crit)
        if self.gpu is not None:
            segmentation_module.cuda()
        segmentation_module.eval()
        self.model = segmentation_module
Esempio n. 8
0
 def __init__(self, encoder_weights, decoder_weights):
     super().__init__()
     self.impl = SegmentationModule(
         net_enc=ModelBuilder.build_encoder(arch="hrnetv2",
                                            fc_dim=720,
                                            weights=encoder_weights),
         net_dec=ModelBuilder.build_decoder(arch="c1",
                                            fc_dim=720,
                                            num_class=150,
                                            weights=decoder_weights,
                                            use_softmax=True),
         crit=None,
     )
Esempio n. 9
0
def segmentation_model_init():
    if not os.path.exists(model_folder):
        os.makedirs(model_folder)

    ENCODER_NAME = 'resnet101'
    DECODER_NAME = 'upernet'
    PRETRAINED_ENCODER_MODEL_URL = 'http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-%s-%s/encoder_epoch_50.pth' % (
        ENCODER_NAME, DECODER_NAME)
    PRETRAINED_DECODER_MODEL_URL = 'http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-%s-%s/decoder_epoch_50.pth' % (
        ENCODER_NAME, DECODER_NAME)

    pretrained_encoder_file = ENCODER_NAME + basename(
        PRETRAINED_ENCODER_MODEL_URL)
    pretrained_decoder_file = DECODER_NAME + basename(
        PRETRAINED_DECODER_MODEL_URL)
    encoder_path = os.path.join(model_folder, pretrained_encoder_file)
    decoder_path = os.path.join(model_folder, pretrained_decoder_file)

    if not os.path.exists(encoder_path):
        urllib.request.urlretrieve(PRETRAINED_ENCODER_MODEL_URL, encoder_path)
    if not os.path.exists(decoder_path):
        urllib.request.urlretrieve(PRETRAINED_DECODER_MODEL_URL, decoder_path)

    # options
    options = SimpleNamespace(fc_dim=2048,
                              num_class=150,
                              imgSizes=[300, 400, 500, 600],
                              imgMaxSize=1000,
                              padding_constant=8,
                              segm_downsampling_rate=8)

    # create model
    builder = ModelBuilder()
    net_encoder = builder.build_encoder(arch=ENCODER_NAME,
                                        weights=encoder_path,
                                        fc_dim=options.fc_dim)
    net_decoder = builder.build_decoder(arch=DECODER_NAME,
                                        weights=decoder_path,
                                        fc_dim=options.fc_dim,
                                        num_class=options.num_class,
                                        use_softmax=True)
    segmentation_module = SegmentationModule(net_encoder, net_decoder,
                                             torch.nn.NLLLoss(ignore_index=-1))
    segmentation_module = segmentation_module.eval()
    torch.set_grad_enabled(False)

    if torch.cuda.is_available():
        segmentation_module = segmentation_module.cuda()
    return segmentation_module, options
Esempio n. 10
0
def load_model_from_cfg(cfg):
    model_config, encoder_path, decoder_path = parse_model_config(cfg)
    net_encoder = ModelBuilder.build_encoder(
        arch=model_config["MODEL"]['arch_encoder'],
        fc_dim=model_config['MODEL']['fc_dim'],
        weights=encoder_path)
    net_decoder = ModelBuilder.build_decoder(
        arch=model_config["MODEL"]['arch_decoder'],
        fc_dim=model_config['MODEL']['fc_dim'],
        num_class=model_config['DATASET']['num_class'],
        weights=decoder_path,
        use_softmax=True)

    crit = torch.nn.NLLLoss(ignore_index=-1)
    segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
    return segmentation_module
Esempio n. 11
0
def load_model():
    net_encoder = ModelBuilder.build_encoder(
    arch='resnet50dilated',
    fc_dim=2048,
    weights='models/encoder_epoch_20.pth')
    net_decoder = ModelBuilder.build_decoder(
    arch='ppm_deepsup',
    fc_dim=2048,
    num_class=150,
    weights='models/decoder_epoch_20.pth',
    use_softmax=True)

    crit = torch.nn.NLLLoss(ignore_index=-1)
    segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
    segmentation_module.eval()
    #segmentation_module.cuda()

    return segmentation_module
Esempio n. 12
0
def semantic_segmentation(image_folder):
    download_segmentation()
    # Network Builders
    net_encoder = ModelBuilder.build_encoder(
        arch='resnet101dilated',
        fc_dim=2048,
        weights='data/segm_data/encoder_epoch_25.pth')
    net_decoder = ModelBuilder.build_decoder(
        arch='ppm_deepsup',
        fc_dim=2048,
        num_class=150,
        weights='data/segm_data/decoder_epoch_25.pth',
        use_softmax=True)

    crit = torch.nn.NLLLoss(ignore_index=-1)
    segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
    segmentation_module.eval()
    segmentation_module.cuda()

    # Load and normalize one image as a singleton tensor batch
    pil_to_tensor = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize(
            mean=[0.485, 0.456, 0.406],  # These are RGB mean+std values
            std=[0.229, 0.224, 0.225])  # across a large photo dataset.
    ])
    pil_image = PIL.Image.open(os.path.join(image_folder,
                                            '000000.jpg')).convert('RGB')
    img_original = np.array(pil_image)
    img_data = pil_to_tensor(pil_image)
    singleton_batch = {'img_data': img_data[None].cuda()}
    output_size = img_data.shape[1:]

    # Run the segmentation at the highest resolution.
    with torch.no_grad():
        scores = segmentation_module(singleton_batch, segSize=output_size)

    # Get the predicted scores for each pixel
    _, pred = torch.max(scores, dim=1)
    pred = pred.cpu()[0].numpy()
    del net_encoder, net_decoder, segmentation_module
    return pred
Esempio n. 13
0
    def __init__(self, model_name, model_path):
        self.model_name = model_name
        self.model_path = model_path
        #
        torch.cuda.set_device(0)

        # Network Builders
        net_encoder = ModelBuilder.build_encoder(arch="hrnetv2",
                                                 fc_dim=720,
                                                 weights=self.model_path)
        net_decoder = ModelBuilder.build_decoder(arch="c1",
                                                 fc_dim=720,
                                                 num_class=2,
                                                 weights=self.model_path,
                                                 use_softmax=True)

        crit = nn.NLLLoss(ignore_index=-1)

        self.segmentation_module = SegmentationModule(net_encoder, net_decoder,
                                                      crit)
Esempio n. 14
0
def main():
    args = build_argparser().parse_args()
    cfg.merge_from_file(args.cfg)
    #cfg.merge_from_list(args.opts)

    logger = setup_logger(distributed_rank=0)   # TODO
    logger.info("Loaded configuration file {}".format(args.cfg))
    logger.info("Running with config:\n{}".format(cfg))

    cfg.MODEL.arch_encoder = cfg.MODEL.arch_encoder.lower()
    cfg.MODEL.arch_decoder = cfg.MODEL.arch_decoder.lower()

    # absolute paths of model weights
    cfg.MODEL.weights_encoder = os.path.join(
        cfg.DIR, 'encoder_' + cfg.TEST.checkpoint)
    cfg.MODEL.weights_decoder = os.path.join(
        cfg.DIR, 'decoder_' + cfg.TEST.checkpoint)

    assert os.path.exists(cfg.MODEL.weights_encoder) and \
        os.path.exists(cfg.MODEL.weights_decoder), "checkpoint does not exitst!"

    # Network Builders
    net_encoder = ModelBuilder.build_encoder(
        arch=cfg.MODEL.arch_encoder,
        fc_dim=cfg.MODEL.fc_dim,
        weights=cfg.MODEL.weights_encoder)
    
    net_decoder = ModelBuilder.build_decoder(
        arch=cfg.MODEL.arch_decoder,
        fc_dim=cfg.MODEL.fc_dim,
        num_class=cfg.DATASET.num_class,
        weights=cfg.MODEL.weights_decoder,
        use_softmax=True)
    
    seg_model = SegmentationModule(net_encoder, net_decoder, None)
    input = torch.randn(1, 3, 320, 320)
    #input_decoder = torch.randn(1, 180, 720, 3, 3)
    segSize = (256, 256)
    #torch.onnx.export(net_encoder, input, "encoder.onnx", opset_version=9, verbose=True)
    torch.onnx.export(seg_model, input, "seg.onnx", opset_version=9, verbose=True)
Esempio n. 15
0
def load_model_from_url(semsegPath):

    model_urls = {
        'encoder':
        'http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth',
        'decoder':
        'http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-resnet50dilated-ppm_deepsup/decoder_epoch_20.pth'
    }

    if not os.path.exists(os.path.join(semsegPath, 'ckpt')):
        os.makedirs(os.path.join(semsegPath, 'ckpt'))

    r = requests.get(model_urls['encoder'])
    with open(os.path.join(semsegPath, 'ckpt/encoder.pth'), 'wb') as f:
        f.write(r.content)
    r = requests.get(model_urls['decoder'])
    with open(os.path.join(semsegPath, 'ckpt/decoder.pth'), 'wb') as f:
        f.write(r.content)

    net_encoder = ModelBuilder.build_encoder(arch='resnet50dilated',
                                             fc_dim=2048,
                                             weights=os.path.join(
                                                 semsegPath,
                                                 'ckpt/encoder.pth'))
    net_decoder = ModelBuilder.build_decoder(arch='ppm_deepsup',
                                             fc_dim=2048,
                                             num_class=150,
                                             weights=os.path.join(
                                                 semsegPath,
                                                 'ckpt/decoder.pth'),
                                             use_softmax=True)

    crit = torch.nn.NLLLoss(ignore_index=-1)
    segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
    segmentation_module.eval()
    segmentation_module.cuda()

    return segmentation_module
Esempio n. 16
0
def create_network(cfg, gpu):
    torch.cuda.set_device(gpu)

    # Network Builders
    net_encoder = ModelBuilder.build_encoder(arch=cfg.MODEL.arch_encoder,
                                             fc_dim=cfg.MODEL.fc_dim,
                                             weights=cfg.MODEL.weights_encoder)
    net_decoder = ModelBuilder.build_decoder(arch=cfg.MODEL.arch_decoder,
                                             fc_dim=cfg.MODEL.fc_dim,
                                             num_class=cfg.DATASET.num_class,
                                             weights=cfg.MODEL.weights_decoder,
                                             use_softmax=True)

    crit = nn.NLLLoss(ignore_index=-1)

    segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)

    segmentation_module.cuda()

    # Main loop
    return segmentation_module, gpu

    print('Inference done!')
def main():
    args = build_argparser().parse_args()
    cfg.merge_from_file(args.cfg)
    #cfg.merge_from_list(args.opts)

    logger = setup_logger(distributed_rank=0)  # TODO
    logger.info("Loaded configuration file {}".format(args.cfg))
    logger.info("Running with config:\n{}".format(cfg))

    cfg.MODEL.arch_encoder = cfg.MODEL.arch_encoder.lower()
    cfg.MODEL.arch_decoder = cfg.MODEL.arch_decoder.lower()

    # absolute paths of model weights
    cfg.MODEL.weights_encoder = os.path.join(cfg.DIR,
                                             'encoder_' + cfg.TEST.checkpoint)
    cfg.MODEL.weights_decoder = os.path.join(cfg.DIR,
                                             'decoder_' + cfg.TEST.checkpoint)

    assert os.path.exists(cfg.MODEL.weights_encoder) and \
        os.path.exists(cfg.MODEL.weights_decoder), "checkpoint does not exitst!"

    # Network Builders
    net_encoder = ModelBuilder.build_encoder(arch=cfg.MODEL.arch_encoder,
                                             fc_dim=cfg.MODEL.fc_dim,
                                             weights=cfg.MODEL.weights_encoder)
    '''
    net_decoder = ModelBuilder.build_decoder(
        arch=cfg.MODEL.arch_decoder,
        fc_dim=cfg.MODEL.fc_dim,
        num_class=cfg.DATASET.num_class,
        weights=cfg.MODEL.weights_decoder,
        use_softmax=True)
    '''
    input = torch.randn(1, 3, 320, 320)
    input_decoder = torch.randn(1, 512, 2048, 40, 40)
    print(input_decoder[-1])
    segSize = (256, 256)
        return [x1, y1, x2, y2]


"""## Loading the segmentation model

Here we load a pretrained segmentation model.  Like any pytorch model, we can call it like a function, or examine the parameters in all the layers.

After loading, we put it on the GPU.  And since we are doing inference, not training, we put the model in eval mode.

# Новый раздел
"""

# Network Builders
net_encoder = ModelBuilder.build_encoder(arch='resnet50dilated',
                                         fc_dim=2048,
                                         weights='ckt/encoder_epoch_20.pth')
net_decoder = ModelBuilder.build_decoder(arch='ppm_deepsup',
                                         fc_dim=2048,
                                         num_class=150,
                                         weights='ckt/decoder_epoch_20.pth',
                                         use_softmax=True)

crit = torch.nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
segmentation_module.eval()
segmentation_module.cuda()
y = Yandex()
"""## Load test data

Now we load and normalize a single test image.  Here we use the commonplace convention of normalizing the image to a scale for which the RGB values of a large photo dataset would have zero mean and unit standard deviation.  (These numbers come from the imagenet dataset.)  With this normalization, the limiiting ranges of RGB values are within about (-2.2 to +2.7).
def huawei_seg(imgs):

    parser = argparse.ArgumentParser(
        description="PyTorch Semantic Segmentation Testing")
    parser.add_argument(
        "--cfg",
        default="config/ade20k-hrnetv2-huawei.yaml",
        metavar="FILE",
        help="path to config file",
        type=str,
    )
    parser.add_argument("--gpu",
                        default=0,
                        type=int,
                        help="gpu id for evaluation")

    args = parser.parse_args()

    cfg.merge_from_file(args.cfg)

    cfg.MODEL.arch_encoder = cfg.MODEL.arch_encoder.lower()
    cfg.MODEL.arch_decoder = cfg.MODEL.arch_decoder.lower()

    # absolute paths of model weights
    cfg.MODEL.weights_encoder = os.path.join(cfg.DIR,
                                             'encoder_' + cfg.TEST.checkpoint)
    cfg.MODEL.weights_decoder = os.path.join(cfg.DIR,
                                             'decoder_' + cfg.TEST.checkpoint)
    #
    imgs = [imgs]
    cfg.list_test = [{'fpath_img': x} for x in imgs]

    torch.cuda.set_device(args.gpu)

    # Network Builders
    net_encoder = ModelBuilder.build_encoder(arch=cfg.MODEL.arch_encoder,
                                             fc_dim=cfg.MODEL.fc_dim,
                                             weights=cfg.MODEL.weights_encoder)
    net_decoder = ModelBuilder.build_decoder(arch=cfg.MODEL.arch_decoder,
                                             fc_dim=cfg.MODEL.fc_dim,
                                             num_class=cfg.DATASET.num_class,
                                             weights=cfg.MODEL.weights_decoder,
                                             use_softmax=True)

    crit = nn.NLLLoss(ignore_index=-1)

    segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)

    # Dataset and Loader
    dataset_test = InferDataset(cfg.list_test, cfg.DATASET)
    loader_test = torch.utils.data.DataLoader(
        dataset_test,
        batch_size=1,
        shuffle=False,
        collate_fn=user_scattered_collate,
        num_workers=5,
        drop_last=True)

    segmentation_module.cuda()
    loader = loader_test
    # Main loop
    segmentation_module.eval()
    pbar = tqdm(total=len(loader))
    for batch_data in loader:
        # process data
        batch_data = batch_data[0]
        segSize = (batch_data['img_ori'].shape[0],
                   batch_data['img_ori'].shape[1])
        img_resized_list = batch_data['img_data']

        with torch.no_grad():
            scores = torch.zeros(1, cfg.DATASET.num_class, segSize[0],
                                 segSize[1])
            scores = async_copy_to(scores, args.gpu)

            for img in img_resized_list:
                feed_dict = batch_data.copy()
                feed_dict['img_data'] = img
                del feed_dict['img_ori']
                del feed_dict['info']
                feed_dict = async_copy_to(feed_dict, args.gpu)

                # forward pass
                pred_tmp = segmentation_module(feed_dict, segSize=segSize)
                # print(pred_tmp.shape)#torch.Size([1, 2, 1024, 1024])
                scores = scores + pred_tmp / len(cfg.DATASET.imgSizes)

            _, pred = torch.max(scores, dim=1)
            pred = as_numpy(pred.squeeze(0).cpu())
        # visualization
        visualize_result((batch_data['img_ori'], batch_data['info']), pred,
                         cfg)
        pbar.update(1)
    #
    return pred
def main(cfg, gpus):
    # Network Builders
    net_encoder = ModelBuilder.build_encoder(
        arch=cfg.MODEL.arch_encoder.lower(),
        fc_dim=cfg.MODEL.fc_dim,
        weights=cfg.MODEL.weights_encoder)
    net_decoder = ModelBuilder.build_decoder(
        arch=cfg.MODEL.arch_decoder.lower(),
        fc_dim=cfg.MODEL.fc_dim,
        num_class=cfg.DATASET.num_class,
        weights=cfg.MODEL.weights_decoder)

    if cfg.MODEL.arch_decoder.endswith('regression'):
        crit = nn.MSELoss(
            reduction="sum"
        )  # Sum for multi-output learning, need to sum across all labels
    else:
        crit = nn.NLLLoss(ignore_index=-1)  # negative log likelihood loss

    if cfg.MODEL.arch_decoder.endswith('deepsup'):
        segmentation_module = SegmentationModule(net_encoder, net_decoder,
                                                 crit, cfg.DATASET.classes,
                                                 cfg.TRAIN.deep_sup_scale)
    else:
        segmentation_module = SegmentationModule(net_encoder, net_decoder,
                                                 crit, cfg.DATASET.classes)

    print("net_encoder")
    print(type(net_encoder))
    print(net_encoder)
    print("net_decoder")
    print(type(net_decoder))
    print(net_decoder)

    # Dataset and Loader
    if cfg.MODEL.arch_decoder.endswith('regression'):
        print("performing regression")
        dataset_train = TrainDatasetRegression(
            cfg.DATASET.root_dataset,
            cfg.DATASET.list_train,
            cfg.DATASET.classes,
            cfg.DATASET,
            batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)

        dataset_val = TrainDatasetRegression(
            cfg.DATASET.root_dataset,
            cfg.DATASET.list_val,
            cfg.DATASET.classes,
            cfg.DATASET,
            batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)
    else:
        dataset_train = TrainDataset(
            cfg.DATASET.root_dataset,
            cfg.DATASET.list_train,
            cfg.DATASET,
            batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)

    loader_train = torch.utils.data.DataLoader(
        dataset_train,
        batch_size=len(gpus),  # we have modified data_parallel
        shuffle=False,  # we do not use this param
        collate_fn=user_scattered_collate,
        num_workers=cfg.TRAIN.workers,
        drop_last=True,
        pin_memory=True)
    loader_val = torch.utils.data.DataLoader(
        dataset_val,
        batch_size=len(gpus),  # we have modified data_parallel
        shuffle=False,  # we do not use this param
        collate_fn=user_scattered_collate,
        num_workers=cfg.TRAIN.workers,
        drop_last=True,
        pin_memory=True)
    print('1 Epoch = {} iters'.format(cfg.TRAIN.epoch_iters))

    # create loader iterator
    iterator_train = iter(loader_train)
    iterator_val = iter(loader_val)

    # load nets into gpu
    if len(gpus) > 1:
        segmentation_module = UserScatteredDataParallel(segmentation_module,
                                                        device_ids=gpus)
        # For sync bn
        patch_replication_callback(segmentation_module)
    segmentation_module.cuda()

    # Set up optimizers
    nets = (net_encoder, net_decoder, crit)
    optimizers = create_optimizers(nets, cfg)

    # Main loop
    history = {
        'train': {
            'epoch': [],
            'loss': []
        },
        'val': {
            'epoch': [],
            'loss': []
        }
    }

    for epoch in range(cfg.TRAIN.start_epoch, cfg.TRAIN.num_epoch):
        train(segmentation_module, iterator_train, optimizers, history,
              epoch + 1, cfg)
        val(segmentation_module, iterator_val, optimizers, history, epoch + 1,
            cfg)

        # checkpointing every 5th epoch
        if (epoch % 5 == 0):
            checkpoint(nets, history, cfg, epoch + 1)

    print('Training Done!')
Esempio n. 21
0
    '''
    net_encoder = ModelBuilder.build_encoder(
        arch='resnet50dilated',
        fc_dim=2048,
        weights='ckpt/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth')
    net_decoder = ModelBuilder.build_decoder(
        arch='ppm_deepsup',
        fc_dim=2048,
        num_class=150,
        weights='ckpt/ade20k-resnet50dilated-ppm_deepsup/decoder_epoch_20.pth',
        use_softmax=True)
    '''
    print("parsing {}".format(args.cfg))
    model_config, encoder_path, decoder_path = parse_model_config(args.cfg)
    net_encoder = ModelBuilder.build_encoder(
        arch = model_config["MODEL"]['arch_encoder'],
        fc_dim = model_config['MODEL']['fc_dim'],
        weights = encoder_path)
    net_decoder = ModelBuilder.build_decoder(
        arch = model_config["MODEL"]['arch_decoder'],
        fc_dim = model_config['MODEL']['fc_dim'],
        num_class = model_config['DATASET']['num_class'],
        weights = decoder_path,
        use_softmax=True)

    crit = torch.nn.NLLLoss(ignore_index=-1)
    segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
    segmentation_module.eval()
    segmentation_module.cuda()

    # creating the videocapture object
    # and reading from the input file
Esempio n. 22
0
# print(torch.load(net_weight,map_location=lambda storage, loc: storage)["encoder"])
#---------------merge model------------------
# dic={}
# net_weight_1='/home/admins/qyl/huawei_compete/semantic-segmentation-pytorch/ckpt/ade20k-hrnetv2-c1/encoder_epoch_20.pth'
# net_weight_2='/home/admins/qyl/huawei_compete/semantic-segmentation-pytorch/ckpt/ade20k-hrnetv2-c1/decoder_epoch_20.pth'
# dic
#
# ["encoder"]=torch.load(net_weight_1)
# dic["decoder"]=torch.load(net_weight_2)
# torch.save(dic,"model_best_20.pth")
# print("merge model finished..........")
#
# Network Builders
model_path = 'model_best.pth'
net_encoder = ModelBuilder.build_encoder(arch="hrnetv2",
                                         fc_dim=720,
                                         weights=model_path)
net_decoder = ModelBuilder.build_decoder(arch="c1",
                                         fc_dim=720,
                                         num_class=2,
                                         weights=model_path,
                                         use_softmax=True)

crit = nn.NLLLoss(ignore_index=-1)

segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
image_demo_dir = "./demo/182_28_9.png"
imags = Image.open(image_demo_dir)
imags = np.array(imags)
imags = imags[512:1024, 0:512, :]
print(imags.shape)
def main(cfg, gpus):
    # Network Builders
    net_encoder = ModelBuilder.build_encoder(
        arch=cfg.MODEL.arch_encoder.lower(),
        fc_dim=cfg.MODEL.fc_dim,
        weights=cfg.MODEL.weights_encoder)
    net_decoder = ModelBuilder.build_decoder(
        arch=cfg.MODEL.arch_decoder.lower(),
        fc_dim=cfg.MODEL.fc_dim,
        num_class=cfg.DATASET.num_class,
        weights=cfg.MODEL.weights_decoder)

    crit = nn.NLLLoss(ignore_index=-1)

    if cfg.MODEL.arch_decoder.endswith('deepsup'):
        segmentation_module = SegmentationModule(net_encoder, net_decoder,
                                                 crit,
                                                 cfg.TRAIN.deep_sup_scale)
    else:
        segmentation_module = SegmentationModule(net_encoder, net_decoder,
                                                 crit)

    # Dataset and Loader
    dataset_train = TrainDataset(cfg.DATASET.root_dataset,
                                 cfg.DATASET.list_train,
                                 cfg.DATASET,
                                 batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)

    loader_train = torch.utils.data.DataLoader(
        dataset_train,
        batch_size=len(gpus),  # we have modified data_parallel
        shuffle=False,  # we do not use this param
        collate_fn=user_scattered_collate,
        num_workers=cfg.TRAIN.workers,
        drop_last=True,
        pin_memory=True)
    print('1 Epoch = {} iters'.format(cfg.TRAIN.epoch_iters))

    # create loader iterator
    iterator_train = iter(loader_train)

    # load nets into gpu
    print(gpus)
    # if len(gpus) > 1:
    if True:
        segmentation_module = UserScatteredDataParallel(segmentation_module,
                                                        device_ids=gpus)
        # For sync bn
        patch_replication_callback(segmentation_module)
    segmentation_module.cuda()

    # Set up optimizers
    nets = (net_encoder, net_decoder, crit)
    optimizers = create_optimizers(nets, cfg)

    # Main loop
    history = {'train': {'epoch': [], 'loss': [], 'acc': []}}

    for epoch in range(cfg.TRAIN.start_epoch, cfg.TRAIN.num_epoch):
        train(segmentation_module, iterator_train, optimizers, history,
              epoch + 1, cfg)

        # checkpointing
        checkpoint(nets, history, cfg, epoch + 1)

    print('Training Done!')
    cfg.DIR, 'encoder_epoch_20.pth')  #'decoder_' + cfg.TEST.checkpoint
cfg.MODEL.weights_decoder = os.path.join(cfg.DIR, 'decoder_epoch_20.pth')

assert os.path.exists(cfg.MODEL.weights_encoder) and \
    os.path.exists(cfg.MODEL.weights_decoder), "checkpoint does not exitst!"

if not os.path.isdir(cfg.TEST.result):
    os.makedirs(cfg.TEST.result)

cfg.TEST.result += 'result'

torch.cuda.set_device(args.gpu)

# Network Builders
net_encoder = ModelBuilder.build_encoder(arch=cfg.MODEL.arch_encoder,
                                         fc_dim=cfg.MODEL.fc_dim,
                                         weights=cfg.MODEL.weights_encoder)
net_decoder = ModelBuilder.build_decoder(arch=cfg.MODEL.arch_decoder,
                                         fc_dim=cfg.MODEL.fc_dim,
                                         num_class=cfg.DATASET.num_class,
                                         weights=cfg.MODEL.weights_decoder,
                                         use_softmax=True)

crit = nn.NLLLoss(ignore_index=-1)

segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)

# generate testing image list
if os.path.isdir(args.imgs):
    imgs = find_recursive(args.imgs)
else:
Esempio n. 25
0
# Network Builders

# upernet resnet101
# net_encoder = ModelBuilder.build_encoder(
#     arch='resnet101',
#     fc_dim=2048,
#     weights='ckpt/ade20k-resnet101-upernet/encoder_epoch_50.pth')
# net_decoder = ModelBuilder.build_decoder(
#     arch='upernet',
#     fc_dim=2048,
#     num_class=150,
#     weights='ckpt/ade20k-resnet101-upernet/decoder_epoch_50.pth',
#     use_softmax=True)
# hrnet
net_encoder = ModelBuilder.build_encoder(
    arch='hrnetv2',
    fc_dim=720,
    weights='../../ckpt/ade20k-hrnetv2-c1/encoder_epoch_30.pth')
net_decoder = ModelBuilder.build_decoder(
    arch='c1',
    fc_dim=720,
    num_class=150,
    weights='../../ckpt/ade20k-hrnetv2-c1/decoder_epoch_30.pth',
    use_softmax=True)

crit = torch.nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
segmentation_module.eval()
segmentation_module.cuda()

pil_to_tensor = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),
Esempio n. 26
0
def inference_prob(
    img,
    device,
    select_model_option="ade20k-resnet50dilated-ppm_deepsup"
):  # select_model_option = "ade20k-mobilenetv2dilated-c1_deepsup" / "ade20k-hrnetv2"
    '''Load the data and preprocess settings
    Input:
        img - the path of our target image
        device - Current device running
        select_model_option - name of NN we use
    '''
    cfg_ss.merge_from_file("ss/config/" + select_model_option + ".yaml")

    logger = setup_logger(distributed_rank=0)  # TODO

    cfg_ss.MODEL.arch_encoder = cfg_ss.MODEL.arch_encoder.lower()
    cfg_ss.MODEL.arch_decoder = cfg_ss.MODEL.arch_decoder.lower()

    # absolute paths of model weights
    cfg_ss.MODEL.weights_encoder = os.path.join(
        'ss/' + cfg_ss.DIR, 'encoder_' + cfg_ss.TEST.checkpoint)
    cfg_ss.MODEL.weights_decoder = os.path.join(
        'ss/' + cfg_ss.DIR, 'decoder_' + cfg_ss.TEST.checkpoint)

    assert os.path.exists(cfg_ss.MODEL.weights_encoder) and os.path.exists(
        cfg_ss.MODEL.weights_decoder), "checkpoint does not exist!"

    # generate testing image list
    imgs = [img]
    assert len(imgs), "imgs should be a path to image (.jpg) or directory."
    cfg_ss.list_test = [{'fpath_img': x} for x in imgs]

    if not os.path.isdir(cfg_ss.TEST.result):
        os.makedirs(cfg_ss.TEST.result)

    if torch.cuda.is_available():
        torch.cuda.set_device(device)

    # Network Builders
    net_encoder = ModelBuilder.build_encoder(
        arch=cfg_ss.MODEL.arch_encoder,
        fc_dim=cfg_ss.MODEL.fc_dim,
        weights=cfg_ss.MODEL.weights_encoder)
    net_decoder = ModelBuilder.build_decoder(
        arch=cfg_ss.MODEL.arch_decoder,
        fc_dim=cfg_ss.MODEL.fc_dim,
        num_class=cfg_ss.DATASET.num_class,
        weights=cfg_ss.MODEL.weights_decoder,
        use_softmax=True)

    crit = nn.NLLLoss(ignore_index=-1)

    segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)

    # Dataset and Loader
    dataset_test = TestDataset(cfg_ss.list_test, cfg_ss.DATASET)
    loader_test = torch.utils.data.DataLoader(
        dataset_test,
        batch_size=cfg_ss.TEST.batch_size,
        shuffle=False,
        collate_fn=user_scattered_collate,
        num_workers=5,
        drop_last=True)

    segmentation_module.to(device)

    # Main loop
    return segmentation_module, loader_test