mmseg返回logits

最新推荐文章于 2024-04-26 20:07:24 发布

carry_hjr

最新推荐文章于 2024-04-26 20:07:24 发布

阅读量3.1k

点赞数

CC 4.0 BY-SA版权

分类专栏：遥感图像处理与深度学习文章标签：深度学习计算机视觉 python

本文链接：https://blog.youkuaiyun.com/carry_hjr/article/details/122551840

遥感图像处理与深度学习专栏收录该内容

22 篇文章

订阅专栏

本文详细介绍了MMSeg中模型的推理过程，包括如何通过不同的配置实现单张图片的预测及多张图片的增强预测，并展示了如何使用API进行实例演示。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

info

主要是为了投票
mmseg/models/segmentors/base.py 里面 forward

        if return_loss:
            return self.forward_train(img, img_metas, **kwargs)
        else:
            return self.forward_test(img, img_metas, **kwargs)

forward_test 里面

        if num_augs == 1:
            return self.simple_test(imgs[0], img_metas[0], **kwargs)
        else:
            #  print('Using aug_test')
            return self.aug_test(imgs, img_metas, **kwargs)

最后是在 mmseg/models/segmentors/encoder_decoder.py 里面实现了 simple_test 和 aug_test ，同时 EncoderDecoder 里面是有test_cfg的那么可以简单传个参数

    def simple_test(self, img, img_meta, rescale=True):
        """Simple test with single image."""
        seg_logit = self.inference(img, img_meta, rescale)

        if self.test_cfg.get("return_logits", False):
            seg_pred = seg_logit
        else:
            seg_pred = seg_logit.argmax(dim=1)

        #  seg_pred = seg_logit.argmax(dim=1)
        if torch.onnx.is_in_onnx_export():
            # our inference backend only support 4D output
            seg_pred = seg_pred.unsqueeze(0)
            return seg_pred
        seg_pred = seg_pred.cpu().numpy()
        # unravel batch dim
        seg_pred = list(seg_pred)
        return seg_pred

效果

from mmseg.apis import init_segmentor, inference_segmentor
config_file = '/home/user/workplace/python/202201kesai_seg/configs/upernet.py'
checkpoint_file = '/home/user/workplace/python/mmsegmentation-0.19.0/work_dirs/2022kesai_seg/upernet/latest.pth'
model = init_segmentor(config_file, checkpoint_file, device='cuda:0')
result = inference_segmentor(model, '/home/user/data/202201_kesai_seg/images/000893_GF.tif')
print(result[0].shape)

(7, 256, 256)

或者用 --out probs.pkl 也可以得到probs